# Buggy code
def get_attributes_from_matrix(X, batch_indices=0, labels=None):
    ne_cells = X.sum(axis=1) > 0
    to_keep = np.where(ne_cells)
    if not ne_cells.all():
        X = X[to_keep]
        removed_idx = np.where(~ne_cells)[0]
        print(
            "Cells with zero expression in all genes considered were removed, the indices of the removed cells "
            "in the expression matrix were:"
        )
        print(removed_idx)
    local_mean, local_var = GeneExpressionDataset.library_size(X)
    batch_indices = (
        batch_indices * np.ones((X.shape[0], 1))
        if type(batch_indices) is int
        else batch_indices[to_keep]
    )
    labels = (
        labels[to_keep].reshape(-1, 1)
        if labels is not None
        else np.zeros_like(batch_indices)
    )
    return X, local_mean, local_var, batch_indices, 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 get_attributes_from_matrix(X, batch_indices=0, labels=None):
    ne_cells = X.sum(axis=1) > 0
    to_keep = np.where(ne_cells)[0]
    if not ne_cells.all():
        X = X[to_keep]
        removed_idx = np.where(~ne_cells)[0]
        print(
            "Cells with zero expression in all genes considered were removed, the indices of the removed cells "
            "in the expression matrix were:"
        )
        print(removed_idx)
    local_mean, local_var = GeneExpressionDataset.library_size(X)
    batch_indices = (
        batch_indices * np.ones((X.shape[0], 1))
        if type(batch_indices) is int
        else batch_indices[to_keep]
    )
    labels = (
        labels[to_keep].reshape(-1, 1)
        if labels is not None
        else np.zeros_like(batch_indices)
    )
    return X, local_mean, local_var, batch_indices, labels