# Buggy code
def extract_data_from_anndata(self, ad: anndata.AnnData):
    data, gene_names, batch_indices, cell_types, labels = None, None, None, None, None
    self.obs = (
        ad.obs
    )  # provide access to observation annotations from the underlying AnnData object.

    if isinstance(ad.X, np.ndarray):
        data = ad.X.copy()  # Dense
    else:
        data = ad.X.toarray()  # Sparse

    select = data.sum(axis=1) > 0  # Take out cells that doesn't express any gene
    data = data[select, :]

    gene_names = np.array(ad.var.index.values, dtype=str)

    if "batch_indices" in self.obs.columns:
        batch_indices = self.obs["batch_indices"].values

    if "cell_types" in self.obs.columns:
        cell_types = self.obs["cell_types"]
        cell_types = cell_types.drop_duplicates().values.astype("str")

    if "labels" in self.obs.columns:
        labels = self.obs["labels"]
    elif "cell_types" in self.obs.columns:
        labels = self.obs["cell_types"].rank(method="dense").values.astype("int")

    return data, gene_names, batch_indices, cell_types, labels


# Traceback (MemoryError)
---------------------------------------------------------------------------
MemoryError                               Traceback (most recent call last)
<ipython-input-3-138de299bfb8> in <module>
1 # load
2 dataset = AnnDataset('data/dataset.h5ad',
----> 3                      save_path=path)

~/anaconda3/envs/pyro/lib/python3.6/site-packages/scvi/dataset/anndata.py in __init__(self, filename, save_path, url, new_n_genes, subset_genes)
36         self.url = url
37
---> 38         data, gene_names = self.download_and_preprocess()
39
40         super().__init__(*GeneExpressionDataset.get_attributes_from_matrix(data),

~/anaconda3/envs/pyro/lib/python3.6/site-packages/scvi/dataset/dataset.py in download_and_preprocess(self)
61     def download_and_preprocess(self):
62         self.download()
---> 63         return self.preprocess()
64
65     def collate_fn(self, batch):

~/anaconda3/envs/pyro/lib/python3.6/site-packages/scvi/dataset/anndata.py in preprocess(self)
52             data = ad.X.copy()  # Dense
53         else:
---> 54             data = ad.X.toarray()  # Sparse
55         select = data.sum(axis=1) > 0  # Take out cells that doesn't express any gene
56         data = data[select, :]

~/anaconda3/envs/pyro/lib/python3.6/site-packages/scipy/sparse/compressed.py in toarray(self, order, out)
960         if out is None and order is None:
961             order = self._swap('cf')[0]
--> 962         out = self._process_toarray_args(order, out)
963         if not (out.flags.c_contiguous or out.flags.f_contiguous):
964             raise ValueError('Output array must be C or F contiguous')

~/anaconda3/envs/pyro/lib/python3.6/site-packages/scipy/sparse/base.py in _process_toarray_args(self, order, out)
1185             return out
1186         else:
-> 1187             return np.zeros(self.shape, dtype=self.dtype, order=order)
1188
1189

MemoryError:

# Fixed code
def extract_data_from_anndata(self, ad: anndata.AnnData):
    data, gene_names, batch_indices, cell_types, labels = None, None, None, None, None
    self.obs = (
        ad.obs
    )  # provide access to observation annotations from the underlying AnnData object.

    # treat all possible cases according to anndata doc
    if isinstance(ad.X, np.ndarray):
        data = ad.X.copy()
    if isinstance(ad.X, pd.DataFrame):
        data = ad.X.values
    if isinstance(ad.X, csr_matrix):
        # keep sparsity above 1 Gb in dense form
        if reduce(operator.mul, ad.X.shape) * ad.X.dtype.itemsize < 1e9:
            data = ad.X.toarray()
        else:
            data = ad.X.copy()

    gene_names = np.array(ad.var.index.values, dtype=str)

    if "batch_indices" in self.obs.columns:
        batch_indices = self.obs["batch_indices"].values

    if "cell_types" in self.obs.columns:
        cell_types = self.obs["cell_types"]
        cell_types = cell_types.drop_duplicates().values.astype("str")

    if "labels" in self.obs.columns:
        labels = self.obs["labels"]
    elif "cell_types" in self.obs.columns:
        labels = self.obs["cell_types"].rank(method="dense").values.astype("int")

    return data, gene_names, batch_indices, cell_types, labels