From dac057862d92e84f84b33c0d69c4cfcafed8864f Mon Sep 17 00:00:00 2001 From: moinfar Date: Mon, 13 Jul 2026 16:01:48 +0200 Subject: [PATCH] Minimal implementation together with tests --- CHANGELOG.md | 5 + src/annbatch/loader.py | 145 ++++++++++++++++++++++--- src/annbatch/types.py | 7 +- tests/test_obsm.py | 233 +++++++++++++++++++++++++++++++++++++++++ 4 files changed, 377 insertions(+), 13 deletions(-) create mode 100644 tests/test_obsm.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 9cb5db31..41124851 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,11 @@ and this project adheres to [Semantic Versioning][]. [keep a changelog]: https://keepachangelog.com/en/1.0.0/ [semantic versioning]: https://semver.org/spec/v2.0.0.html +## [Unreleased] + +### Feature +- Add an `obsm_keys` argument to {meth}`annbatch.Loader.add_adata`/{meth}`annbatch.Loader.add_adatas` (and an `obsm` argument to {meth}`annbatch.Loader.add_dataset`/{meth}`annbatch.Loader.add_datasets`) to load dense {attr}`~anndata.AnnData.obsm` arrays (e.g. a scVI embedding `"X_emb"`) alongside `X`. Requested arrays are fetched concurrently with `X` and yielded, row-aligned, under a new `obsm` key of {class}`~annbatch.types.LoaderOutput` (`None` when no keys are requested). + ## [0.2.1] ### Feature diff --git a/src/annbatch/loader.py b/src/annbatch/loader.py index e5d619b3..92747c76 100644 --- a/src/annbatch/loader.py +++ b/src/annbatch/loader.py @@ -179,6 +179,7 @@ class Loader[ _train_datasets: list[BackingArray] _obs: list[pd.DataFrame] | None = None + _obsm: list[dict[str, ZarrArray | np.ndarray]] | None = None _var: pd.DataFrame | None = None _return_index: bool = False _shapes: list[tuple[int, int]] @@ -384,6 +385,8 @@ def use_collection( def add_adatas( self, adatas: list[ad.AnnData], + *, + obsm_keys: list[str] | None = None, ) -> Self: """Append adatas to this dataset. @@ -391,14 +394,17 @@ def add_adatas( ---------- adatas List of :class:`anndata.AnnData` objects, with :class:`zarr.Array`, :class:`scipy.sparse.csr_matrix`, :class:`scipy.sparse.csr_array`, :class:`numpy.ndarray`, or :class:`anndata.abc.CSRDataset` as the data matrix in :attr:`~anndata.AnnData.X`, and :attr:`~anndata.AnnData.obs` containing annotations to yield in a :class:`pandas.DataFrame`. + obsm_keys + Keys of :attr:`~anndata.AnnData.obsm` (e.g. a pretrained-model embedding ``"X_emb"``) to load and yield alongside ``X``. + Each referenced array must be dense (a :class:`numpy.ndarray` or :class:`zarr.Array`) and present in every added adata. """ check_lt_1([len(adatas)], ["Number of adatas"]) for adata in adatas: - dataset, obs, var = self._prepare_dataset_obs_and_var(adata) - self._add_dataset_unchecked(dataset, obs, var) + dataset, obs, var, obsm = self._prepare_dataset_obs_and_var(adata, obsm_keys) + self._add_dataset_unchecked(dataset, obs, var, obsm) return self - def add_adata(self, adata: ad.AnnData) -> Self: + def add_adata(self, adata: ad.AnnData, *, obsm_keys: list[str] | None = None) -> Self: """Append an adata to this dataset. Parameters @@ -406,14 +412,17 @@ def add_adata(self, adata: ad.AnnData) -> Self: adata A :class:`anndata.AnnData` object, with :class:`zarr.Array`, :class:`scipy.sparse.csr_matrix`, :class:`scipy.sparse.csr_array`, :class:`numpy.ndarray`, or :class:`anndata.abc.CSRDataset` as the data matrix in :attr:`~anndata.AnnData.X`, and :attr:`~anndata.AnnData.obs` containing annotations to yield in a :class:`pandas.DataFrame`. :attr:`~anndata.AnnData.var` must match the ``var`` of any previously added datasets. + obsm_keys + Keys of :attr:`~anndata.AnnData.obsm` (e.g. a pretrained-model embedding ``"X_emb"``) to load and yield alongside ``X``. + Each referenced array must be dense (a :class:`numpy.ndarray` or :class:`zarr.Array`) and match the ``obsm`` keys of any previously added datasets. """ - dataset, obs, var = self._prepare_dataset_obs_and_var(adata) - self.add_dataset(dataset, obs, var) + dataset, obs, var, obsm = self._prepare_dataset_obs_and_var(adata, obsm_keys) + self.add_dataset(dataset, obs, var, obsm) return self def _prepare_dataset_obs_and_var( - self, adata: ad.AnnData - ) -> tuple[BackingArray, pd.DataFrame | None, pd.DataFrame | None]: + self, adata: ad.AnnData, obsm_keys: list[str] | None = None + ) -> tuple[BackingArray, pd.DataFrame | None, pd.DataFrame | None, dict[str, ZarrArray | np.ndarray] | None]: dataset = adata.X obs = adata.obs var = adata.var @@ -421,8 +430,14 @@ def _prepare_dataset_obs_and_var( obs = None if not isinstance(dataset, BackingArray_T.__value__): raise TypeError(f"Found {type(dataset)} but only {BackingArray_T.__value__} are usable") + obsm = None + if obsm_keys is not None: + missing = [k for k in obsm_keys if k not in adata.obsm] + if missing: + raise KeyError(f"obsm keys {missing} not found in adata.obsm (available: {list(adata.obsm)})") + obsm = {k: adata.obsm[k] for k in obsm_keys} - return cast("BackingArray", dataset), obs, var + return cast("BackingArray", dataset), obs, var, obsm @validate_sampler def add_datasets( @@ -430,6 +445,7 @@ def add_datasets( datasets: list[BackingArray], obs: list[pd.DataFrame] | None = None, var: list[pd.DataFrame] | None = None, + obsm: list[dict[str, ZarrArray | np.ndarray]] | None = None, ) -> Self: """Append datasets to this dataset. @@ -443,13 +459,19 @@ def add_datasets( var List of :class:`~pandas.DataFrame` for annotating features, generally from :attr:`anndata.AnnData.var`. All var DataFrames must be identical. + obsm + List of dicts mapping an obsm key (e.g. an embedding ``"X_emb"``) to a dense + :class:`numpy.ndarray` or :class:`zarr.Array`, generally from :attr:`anndata.AnnData.obsm`. + All dicts must share the same keys, and each array's first axis must match the corresponding dataset's. """ if obs is None: obs = [None] * len(datasets) if var is None: var = [None] * len(datasets) - for ds, o, v in zip(datasets, obs, var, strict=True): - self._add_dataset_unchecked(ds, o, v) + if obsm is None: + obsm = [None] * len(datasets) + for ds, o, v, m in zip(datasets, obs, var, obsm, strict=True): + self._add_dataset_unchecked(ds, o, v, m) return self @validate_sampler @@ -458,6 +480,7 @@ def add_dataset( dataset: BackingArray, obs: pd.DataFrame | None = None, var: pd.DataFrame | None = None, + obsm: dict[str, ZarrArray | np.ndarray] | None = None, ) -> Self: """Append a dataset to this dataset. @@ -470,8 +493,12 @@ def add_dataset( var :class:`~pandas.DataFrame` var, generally from :attr:`anndata.AnnData.var`. :attr:`~anndata.AnnData.var` must match the ``var`` of any previously added datasets. + obsm + Dict mapping an obsm key (e.g. an embedding ``"X_emb"``) to a dense :class:`numpy.ndarray` + or :class:`zarr.Array`, generally from :attr:`anndata.AnnData.obsm`. + Its keys must match the ``obsm`` of any previously added datasets, and each array's first axis must match ``dataset``. """ - self._add_dataset_unchecked(dataset, obs, var) + self._add_dataset_unchecked(dataset, obs, var, obsm) return self def _add_dataset_unchecked( @@ -479,6 +506,7 @@ def _add_dataset_unchecked( dataset: BackingArray, obs: pd.DataFrame | None = None, var: pd.DataFrame | None = None, + obsm: dict[str, ZarrArray | np.ndarray] | None = None, ) -> Self: if len(self._train_datasets) > 0: if self._obs is None and obs is not None: @@ -497,6 +525,18 @@ def _add_dataset_unchecked( raise ValueError( "Cannot add a dataset without var when training datasets have already been added with var" ) + if self._obsm is None and obsm is not None: + raise ValueError( + "Cannot add a dataset with obsm when training datasets have already been added without obsm" + ) + if self._obsm is not None and obsm is None: + raise ValueError( + "Cannot add a dataset without obsm when training datasets have already been added with obsm" + ) + if self._obsm is not None and obsm is not None and set(obsm) != set(self._obsm[0]): + raise ValueError( + f"All datasets must have identical obsm keys. Expected {sorted(self._obsm[0])} but got {sorted(obsm)}." + ) if not isinstance(dataset, self.dataset_type): raise ValueError( f"All datasets on a given loader must be of the same type {self.dataset_type} but got {type(dataset)}" @@ -513,6 +553,28 @@ def _add_dataset_unchecked( raise TypeError("obs must be a pandas DataFrame") if not isinstance(var, pd.DataFrame) and var is not None: raise TypeError("var must be a pandas DataFrame") + if obsm is not None: + for key, arr in obsm.items(): + if not isinstance(arr, ZarrArray | np.ndarray): + raise TypeError( + f"obsm[{key!r}] must be a dense numpy.ndarray or zarr.Array, got {type(arr)}. " + "Sparse obsm is not supported." + ) + if arr.shape[0] != dataset.shape[0]: + raise ValueError( + f"obsm[{key!r}] has {arr.shape[0]} rows but the dataset has {dataset.shape[0]} observations." + ) + if self._obsm is not None: + existing = self._obsm[0][key] + if arr.shape[1:] != existing.shape[1:]: + raise ValueError( + f"obsm[{key!r}] feature shape {arr.shape[1:]} does not match the " + f"existing shape {existing.shape[1:]}." + ) + if arr.dtype != existing.dtype: + raise ValueError( + f"obsm[{key!r}] dtype {arr.dtype!r} does not match the existing dtype {existing.dtype!r}." + ) datasets = self._train_datasets + [dataset] check_var_shapes(datasets) self._dtypes_homogeneous = self._datasets_share_dtype(datasets) @@ -529,6 +591,10 @@ def _add_dataset_unchecked( self._obs += [obs] elif obs is not None: # obs dont exist yet, but are being added for the first time self._obs = [obs] + if self._obsm is not None: # obsm exist + self._obsm += [obsm] + elif obsm is not None: # obsm dont exist yet, but are being added for the first time + self._obsm = [obsm] # var is the same across all datasets (describes variables/features) if self._var is None and var is not None: self._var = var @@ -968,6 +1034,48 @@ async def _index_datasets( return out + async def _index_obsm( + self, + dataset_index_to_rows: OrderedDict[int, np.ndarray], + ) -> dict[str, np.ndarray]: + """Fetch each requested obsm array into a contiguous buffer, in the same dataset order as ``X``. + + obsm arrays are always dense, so a single buffer per key is preallocated (sized to the total + number of requested rows) and filled by concurrent per-dataset fetches, mirroring the dense + path of :meth:`_index_datasets`. The row order matches that of the ``X`` buffer, so the same + ``inv``/split indexing applies downstream. + """ + keys = list(self._obsm[0]) + total_rows = sum(len(rows) for rows in dataset_index_to_rows.values()) + out: dict[str, np.ndarray] = {} + tasks = [] + for key in keys: + first = self._obsm[next(iter(dataset_index_to_rows))][key] + buffer = self._alloc((total_rows, *first.shape[1:]), first.dtype, use_pinned=self._preload_to_gpu) + out[key] = buffer + row_offset = 0 + for dataset_idx, rows in dataset_index_to_rows.items(): + nrows = len(rows) + tasks.append( + self._fetch_data(self._obsm[dataset_idx][key], rows, buffer[row_offset : row_offset + nrows]) + ) + row_offset += nrows + await asyncio.gather(*tasks) + return out + + async def _index_x_and_obsm( + self, + dataset_index_to_rows: OrderedDict[int, np.ndarray], + ) -> tuple[CSRContainer | np.ndarray, dict[str, np.ndarray] | None]: + """Fetch ``X`` and (if requested) the obsm arrays concurrently in a single event loop.""" + if self._obsm is None: + return await self._index_datasets(dataset_index_to_rows), None + x, obsm = await asyncio.gather( + self._index_datasets(dataset_index_to_rows), + self._index_obsm(dataset_index_to_rows), + ) + return x, obsm + def __iter__( self, ) -> Iterator[LoaderOutput[OutputInMemoryArray]]: @@ -1017,7 +1125,9 @@ def __iter__( inv = inv_buffer[:n] inv = positions[order] - raw_out: CSRContainer | np.ndarray = zsync.sync(self._index_datasets(dataset_index_to_rows)) + raw_out: CSRContainer | np.ndarray + raw_obsm: dict[str, np.ndarray] | None + raw_out, raw_obsm = zsync.sync(self._index_x_and_obsm(dataset_index_to_rows)) if is_sparse: in_memory_data = self._sp_module.csr_matrix( @@ -1028,15 +1138,26 @@ def __iter__( else: in_memory_data = self._np_module.asarray(raw_out) + in_memory_obsm: dict[str, np.ndarray] | None = ( + {key: self._np_module.asarray(arr) for key, arr in raw_obsm.items()} if raw_obsm is not None else None + ) + concatenated_obs: None | pd.DataFrame = self._maybe_accumulate_obs(dataset_index_to_rows) in_memory_indices: None | np.ndarray = self._maybe_accumulate_indices(dataset_index_to_rows) for split in splits: sel = inv[split] data = in_memory_data[sel] + obsm_out: dict[str, OutputInMemoryArray_T] | None = None + if in_memory_obsm is not None: + obsm_out = { + key: arr[sel] if self._to is None else convert(arr[sel], self._preload_to_gpu, self._to) + for key, arr in in_memory_obsm.items() + } yield { "X": data if self._to is None else convert(data, self._preload_to_gpu, self._to), "obs": concatenated_obs.iloc[sel] if concatenated_obs is not None else None, "var": self._var, + "obsm": obsm_out, "index": in_memory_indices[sel] if in_memory_indices is not None else None, } diff --git a/src/annbatch/types.py b/src/annbatch/types.py index cdd28082..f0721041 100644 --- a/src/annbatch/types.py +++ b/src/annbatch/types.py @@ -52,9 +52,14 @@ class LoadRequest(TypedDict): class LoaderOutput[OutputInMemoryArray: OutputInMemoryArray_T](TypedDict): - """The output of the loader, the "data matrix" with its obs, optional, var, optional, and index, also optional.""" + """The output of the loader: the "data matrix" ``X`` with its ``obs``, ``var``, ``obsm`` and ``index`` (all optional). + + ``obsm`` is ``None`` unless obsm keys were requested (see :meth:`~annbatch.Loader.add_adatas`); when present it is a + dict mapping each requested key to a dense per-batch array row-aligned with ``X``. + """ X: OutputInMemoryArray_T.__value__ # TODO: remove after sphinx 9 - myst compat obs: pd.DataFrame | None var: pd.DataFrame | None + obsm: dict[str, OutputInMemoryArray_T.__value__] | None index: np.ndarray | None diff --git a/tests/test_obsm.py b/tests/test_obsm.py new file mode 100644 index 00000000..d6ebbb58 --- /dev/null +++ b/tests/test_obsm.py @@ -0,0 +1,233 @@ +from __future__ import annotations + +from importlib.util import find_spec +from typing import TYPE_CHECKING + +import anndata as ad +import numpy as np +import pandas as pd +import pytest +import scipy.sparse as sp +import zarr + +from annbatch import Loader + +if TYPE_CHECKING: + from pathlib import Path + +skip_if_no_numba = pytest.mark.skipif( + find_spec("numba") is None, reason="Can't test for in-memory sparse without numba" +) + +N_VAR = 12 +N_EMB = 5 +SIZES = (50, 30, 40) + + +def _build_adatas( + rng: np.random.Generator, *, sparse: bool +) -> tuple[list[ad.AnnData], np.ndarray, np.ndarray, np.ndarray]: + """Build imaginary adatas with X, obsm['X_emb'], obs (batch+label) and var. + + Returns the adatas plus the globally-concatenated X, X_emb and label arrays so + that batches can be checked against the original rows via ``return_index``. + """ + adatas, x_all, emb_all, label_all = [], [], [], [] + for k, n in enumerate(SIZES): + x = rng.random((n, N_VAR)).astype("f4") + emb = rng.random((n, N_EMB)).astype("f4") + labels = rng.integers(0, 4, size=n) + adatas.append( + ad.AnnData( + X=sp.csr_matrix(x) if sparse else x, + obs=pd.DataFrame( + {"batch": pd.Categorical([f"donor{k}"] * n), "label": labels}, + index=[f"donor{k}_cell_{i}" for i in range(n)], + ), + var=pd.DataFrame(index=[f"gene_{i}" for i in range(N_VAR)]), + obsm={"X_emb": emb}, + ) + ) + x_all.append(x) + emb_all.append(emb) + label_all.append(labels) + return adatas, np.vstack(x_all), np.vstack(emb_all), np.concatenate(label_all) + + +def _to_dense_np(x) -> np.ndarray: + return x.toarray() if sp.issparse(x) else np.asarray(x) + + +@pytest.mark.parametrize("shuffle", [True, False], ids=["shuffled", "unshuffled"]) +@pytest.mark.parametrize( + "sparse", + [pytest.param(False, id="dense"), pytest.param(True, id="sparse", marks=skip_if_no_numba)], +) +def test_obsm_in_memory_alignment(*, sparse: bool, shuffle: bool): + """obsm['X_emb'] is yielded and stays row-aligned with X, obs and index.""" + rng = np.random.default_rng(0) + adatas, x_all, emb_all, label_all = _build_adatas(rng, sparse=sparse) + + loader = Loader( + batch_size=16, + chunk_size=4, + preload_nchunks=8, + shuffle=shuffle, + return_index=True, + to=None, + preload_to_gpu=False, + rng=np.random.default_rng(1), + ).add_adatas(adatas, obsm_keys=["X_emb"]) + + seen_idx = [] + for batch in loader: + assert "obsm" in batch + emb = np.asarray(batch["obsm"]["X_emb"]) + idx = batch["index"] + x = _to_dense_np(batch["X"]) + # shapes line up across X, obsm and obs + assert emb.shape == (len(idx), N_EMB) + assert x.shape[0] == emb.shape[0] == batch["obs"].shape[0] + # every row matches the original global rows at these indices + np.testing.assert_allclose(emb, emb_all[idx]) + np.testing.assert_allclose(x, x_all[idx], atol=1e-6) + np.testing.assert_array_equal(batch["obs"]["label"].to_numpy(), label_all[idx]) + seen_idx.append(idx) + + all_idx = np.concatenate(seen_idx) + assert sorted(all_idx.tolist()) == list(range(sum(SIZES))) + + +def test_obsm_multiple_keys(): + """Several obsm keys can be requested at once and are all yielded.""" + rng = np.random.default_rng(3) + adatas = [] + for n in SIZES: + adatas.append( + ad.AnnData( + X=rng.random((n, N_VAR)).astype("f4"), + obs=pd.DataFrame({"label": rng.integers(0, 3, n)}), + var=pd.DataFrame(index=[f"gene_{i}" for i in range(N_VAR)]), + obsm={"X_emb": rng.random((n, N_EMB)).astype("f4"), "X_pca": rng.random((n, 3)).astype("f4")}, + ) + ) + loader = Loader( + batch_size=16, chunk_size=4, preload_nchunks=8, shuffle=True, to=None, preload_to_gpu=False + ).add_adatas(adatas, obsm_keys=["X_emb", "X_pca"]) + batch = next(iter(loader)) + assert set(batch["obsm"]) == {"X_emb", "X_pca"} + assert batch["obsm"]["X_emb"].shape[1] == N_EMB + assert batch["obsm"]["X_pca"].shape[1] == 3 + + +def test_obsm_none_by_default(): + """Without ``obsm_keys`` the batch carries an explicit ``obsm=None``.""" + rng = np.random.default_rng(4) + adata = ad.AnnData( + X=rng.random((20, N_VAR)).astype("f4"), + obs=pd.DataFrame({"label": rng.integers(0, 3, 20)}), + var=pd.DataFrame(index=[f"gene_{i}" for i in range(N_VAR)]), + obsm={"X_emb": rng.random((20, N_EMB)).astype("f4")}, + ) + loader = Loader(batch_size=8, chunk_size=2, preload_nchunks=4, to=None, preload_to_gpu=False).add_adata(adata) + batch = next(iter(loader)) + assert batch["obsm"] is None + + +def test_obsm_on_disk_alignment(tmp_path: Path): + """obsm backed by an on-disk zarr.Array is fetched and stays aligned.""" + rng = np.random.default_rng(5) + adatas_mem, x_all, emb_all, _ = _build_adatas(rng, sparse=False) + for k, a in enumerate(adatas_mem): + a.write_zarr(tmp_path / f"a{k}.zarr") + + def load(k: int) -> ad.AnnData: + g = zarr.open_group(tmp_path / f"a{k}.zarr", mode="r") + return ad.AnnData( + X=g["X"], + obs=ad.io.read_elem(g["obs"]), + var=pd.DataFrame(index=pd.Index(ad.io.read_elem(g["var"][g["var"].attrs["_index"]]))), + obsm={"X_emb": g["obsm"]["X_emb"]}, + ) + + adatas = [load(k) for k in range(len(SIZES))] + assert isinstance(adatas[0].obsm["X_emb"], zarr.Array) + + loader = Loader( + batch_size=6, + chunk_size=3, + preload_nchunks=4, + shuffle=True, + return_index=True, + to=None, + preload_to_gpu=False, + rng=np.random.default_rng(6), + ).add_adatas(adatas, obsm_keys=["X_emb"]) + + seen = 0 + for batch in loader: + idx = batch["index"] + np.testing.assert_allclose(np.asarray(batch["obsm"]["X_emb"]), emb_all[idx]) + np.testing.assert_allclose(_to_dense_np(batch["X"]), x_all[idx], atol=1e-6) + seen += len(idx) + assert seen == sum(SIZES) + + +def test_obsm_missing_key_raises(): + rng = np.random.default_rng(7) + adata = ad.AnnData( + X=rng.random((10, N_VAR)).astype("f4"), + var=pd.DataFrame(index=[f"gene_{i}" for i in range(N_VAR)]), + obsm={"X_emb": rng.random((10, N_EMB)).astype("f4")}, + ) + with pytest.raises(KeyError, match="not found in adata.obsm"): + Loader(batch_size=2, chunk_size=2, preload_nchunks=2, to=None, preload_to_gpu=False).add_adata( + adata, obsm_keys=["does_not_exist"] + ) + + +def test_obsm_inconsistent_presence_raises(): + rng = np.random.default_rng(8) + + def mk(*, with_obsm: bool) -> ad.AnnData: + return ad.AnnData( + X=rng.random((10, N_VAR)).astype("f4"), + obs=pd.DataFrame({"label": np.zeros(10)}), + var=pd.DataFrame(index=[f"gene_{i}" for i in range(N_VAR)]), + obsm={"X_emb": rng.random((10, N_EMB)).astype("f4")} if with_obsm else {}, + ) + + loader = Loader(batch_size=2, chunk_size=2, preload_nchunks=2, to=None, preload_to_gpu=False) + loader.add_adata(mk(with_obsm=False)) + with pytest.raises(ValueError, match="without obsm"): + loader.add_adata(mk(with_obsm=True), obsm_keys=["X_emb"]) + + +def test_obsm_mismatched_feature_shape_raises(): + rng = np.random.default_rng(9) + a1 = ad.AnnData( + X=rng.random((10, N_VAR)).astype("f4"), + var=pd.DataFrame(index=[f"gene_{i}" for i in range(N_VAR)]), + obsm={"X_emb": rng.random((10, N_EMB)).astype("f4")}, + ) + a2 = ad.AnnData( + X=rng.random((10, N_VAR)).astype("f4"), + var=pd.DataFrame(index=[f"gene_{i}" for i in range(N_VAR)]), + obsm={"X_emb": rng.random((10, N_EMB + 3)).astype("f4")}, + ) + with pytest.raises(ValueError, match="feature shape"): + Loader(batch_size=2, chunk_size=2, preload_nchunks=2, to=None, preload_to_gpu=False).add_adatas( + [a1, a2], obsm_keys=["X_emb"] + ) + + +def test_obsm_wrong_nobs_raises(): + rng = np.random.default_rng(10) + adata = ad.AnnData( + X=rng.random((10, N_VAR)).astype("f4"), + var=pd.DataFrame(index=[f"gene_{i}" for i in range(N_VAR)]), + ) + with pytest.raises(ValueError, match="rows but the dataset has"): + Loader(batch_size=2, chunk_size=2, preload_nchunks=2, to=None, preload_to_gpu=False).add_dataset( + adata.X, obsm={"X_emb": rng.random((7, N_EMB)).astype("f4")} + )