# Buggy code
def get_normalized_expression(
    self,
    adata: Optional[AnnData] = None,
    indices: Optional[Sequence[int]] = None,
    transform_batch: Optional[Sequence[Union[Number, str]]] = None,
    gene_list: Optional[Sequence[str]] = None,
    library_size: Union[float, Literal["latent"]] = 1,
    n_samples: int = 1,
    batch_size: Optional[int] = None,
    return_mean: bool = True,
    return_numpy: Optional[bool] = None,
) -> Union[np.ndarray, pd.DataFrame]:
    r"""
    Returns the normalized (decoded) gene expression.

    This is denoted as :math:`\rho_n` in the scVI paper.

    Parameters
    ----------
    adata
        AnnData object with equivalent structure to initial AnnData. If `None`, defaults to the
        AnnData object used to initialize the model.
    indices
        Indices of cells in adata to use. If `None`, all cells are used.
    transform_batch
        Batch to condition on.
        If transform_batch is:

        - None, then real observed batch is used.
        - int, then batch transform_batch is used.
    gene_list
        Return frequencies of expression for a subset of genes.
        This can save memory when working with large datasets and few genes are
        of interest.
    library_size
        Scale the expression frequencies to a common library size.
        This allows gene expression levels to be interpreted on a common scale of relevant
        magnitude. If set to `"latent"`, use the latent libary size.
    n_samples
        Number of posterior samples to use for estimation.
    batch_size
        Minibatch size for data loading into model. Defaults to `scvi.settings.batch_size`.
    return_mean
        Whether to return the mean of the samples.
    return_numpy
        Return a :class:`~numpy.ndarray` instead of a :class:`~pandas.DataFrame`. DataFrame includes
        gene names as columns. If either `n_samples=1` or `return_mean=True`, defaults to `False`.
        Otherwise, it defaults to `True`.

    Returns
    -------
    If `n_samples` > 1 and `return_mean` is False, then the shape is `(samples, cells, genes)`.
    Otherwise, shape is `(cells, genes)`. In this case, return type is :class:`~pandas.DataFrame` unless `return_numpy` is True.
    """
    adata = self._validate_anndata(adata)
    scdl = self._make_scvi_dl(adata=adata, indices=indices, batch_size=batch_size)

    if not isinstance(transform_batch, IterableClass):
        transform_batch = [transform_batch]

    transform_batch = _get_batch_code_from_category(adata, transform_batch)

    if gene_list is None:
        gene_mask = slice(None)
    else:
        all_genes = _get_var_names_from_setup_anndata(adata)
        gene_mask = [True if gene in gene_list else False for gene in all_genes]

    if n_samples > 1 and return_mean is False:
        if return_numpy is False:
            logger.warning(
                "return_numpy must be True if n_samples > 1 and return_mean is False, returning np.ndarray"
            )
        return_numpy = True
    if indices is None:
        indices = np.arange(adata.n_obs)
    if library_size == "latent":
        model_fn = self.model.get_sample_rate
        scaling = 1
    else:
        model_fn = self.model.get_sample_scale
        scaling = library_size

    exprs = []
    for tensors in scdl:
        x = tensors[_CONSTANTS.X_KEY]
        batch_idx = tensors[_CONSTANTS.BATCH_KEY]
        labels = tensors[_CONSTANTS.LABELS_KEY]
        per_batch_exprs = []
        for batch in transform_batch:
            output = model_fn(
                x,
                batch_index=batch_idx,
                y=labels,
                n_samples=n_samples,
                transform_batch=batch,
            )[..., gene_mask]
            output *= scaling
            output = output.cpu().numpy()
            per_batch_exprs.append(output)
        per_batch_exprs = np.stack(
            per_batch_exprs
        )  # shape is (len(transform_batch) x batch_size x n_var)
        exprs += [per_batch_exprs.mean(0)]

    if n_samples > 1:
        # The -2 axis correspond to cells.
        exprs = np.concatenate(exprs, axis=-2)
    else:
        exprs = np.concatenate(exprs, axis=0)
    if n_samples > 1 and return_mean:
        exprs = exprs.mean(0)

    if return_numpy is None or return_numpy is False:
        return pd.DataFrame(
            exprs,
            columns=adata.var_names[gene_mask],
            index=adata.obs_names[indices],
        )
    else:
        return exprs


# Traceback (ValueError)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-143-b1df721788f5> in <module>
2 for tissue in np.unique(adata.obs['tissue']):
3     sub_adata = adata[adata.obs['tissue']==tissue]
----> 4     de_celltype[tissue] = vae.differential_expression(sub_adata, groupby = 'Propagated.Annotation', batch_correction=True)

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/scvi/core/models/rnamixin.py in differential_expression(self, adata, groupby, group1, group2, idx1, idx2, mode, delta, batch_size, all_stats, batch_correction, batchid1, batchid2, **kwargs)
187             batch_size=batch_size,
188         )
--> 189         result = _de_core(
190             adata,
191             model_fn,

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/scvi/core/models/_utils.py in _de_core(adata, model_fn, groupby, group1, group2, idx1, idx2, all_stats, all_stats_fn, col_names, mode, batchid1, batchid2, delta, batch_correction, **kwargs)
59             cell_idx2 = (adata.obs[groupby] == group2).to_numpy().ravel()
60
---> 61         all_info = dc.get_bayes_factors(
62             cell_idx1,
63             cell_idx2,

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/scvi/core/utils/differential.py in get_bayes_factors(self, idx1, idx2, mode, batchid1, batchid2, use_observed_batches, n_samples, use_permutation, m_permutation, change_fn, m1_domain_fn, delta, cred_interval_lvls)
168         eps = 1e-8  # used for numerical stability
169         # Normalized means sampling for both populations
--> 170         scales_batches_1 = self.scale_sampler(
171             selection=idx1,
172             batchid=batchid1,

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/torch/autograd/grad_mode.py in decorate_no_grad(*args, **kwargs)
47         def decorate_no_grad(*args, **kwargs):
48             with self:
---> 49                 return func(*args, **kwargs)
50         return decorate_no_grad
51

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/scvi/core/utils/differential.py in scale_sampler(self, selection, n_samples, n_samples_per_cell, batchid, use_observed_batches, give_mean)
392             idx = np.random.choice(np.arange(self.adata.shape[0])[selection], n_samples)
393             px_scales.append(
--> 394                 self.model_fn(self.adata, indices=idx, transform_batch=batch_idx)
395             )
396             batch_idx = batch_idx if batch_idx is not None else np.nan

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/torch/autograd/grad_mode.py in decorate_no_grad(*args, **kwargs)
47         def decorate_no_grad(*args, **kwargs):
48             with self:
---> 49                 return func(*args, **kwargs)
50         return decorate_no_grad
51

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/scvi/core/models/rnamixin.py in get_normalized_expression(self, adata, indices, transform_batch, gene_list, library_size, n_samples, batch_size, return_mean, return_numpy)
82         scdl = self._make_scvi_dl(adata=adata, indices=indices, batch_size=batch_size)
83         if transform_batch is not None:
---> 84             transform_batch = _get_batch_code_from_category(adata, transform_batch)
85
86         if gene_list is None:

/data/yosef2/users/chenling/miniconda3/envs/scvi-tools/lib/python3.8/site-packages/scvi/model/_utils.py in _get_batch_code_from_category(adata, category)
131     batch_mappings = categorical_mappings["_scvi_batch"]["mapping"]
132     if category not in batch_mappings:
--> 133         raise ValueError('"{}" not a valid batch category.'.format(category))
134     return np.where(batch_mappings == category)[0][0]

ValueError: "0" not a valid batch category.```

#### Versions:
<!-- Output of scvi.__version__ -->
VERSION

<!-- Relevant screenshots -->

# Fixed code
def get_normalized_expression(
    self,
    adata: Optional[AnnData] = None,
    indices: Optional[Sequence[int]] = None,
    transform_batch: Optional[Sequence[Union[Number, str]]] = None,
    gene_list: Optional[Sequence[str]] = None,
    library_size: Union[float, Literal["latent"]] = 1,
    n_samples: int = 1,
    batch_size: Optional[int] = None,
    return_mean: bool = True,
    return_numpy: Optional[bool] = None,
) -> Union[np.ndarray, pd.DataFrame]:
    r"""
    Returns the normalized (decoded) gene expression.

    This is denoted as :math:`\rho_n` in the scVI paper.

    Parameters
    ----------
    adata
        AnnData object with equivalent structure to initial AnnData. If `None`, defaults to the
        AnnData object used to initialize the model.
    indices
        Indices of cells in adata to use. If `None`, all cells are used.
    transform_batch
        Batch to condition on.
        If transform_batch is:

        - None, then real observed batch is used.
        - int, then batch transform_batch is used.
    gene_list
        Return frequencies of expression for a subset of genes.
        This can save memory when working with large datasets and few genes are
        of interest.
    library_size
        Scale the expression frequencies to a common library size.
        This allows gene expression levels to be interpreted on a common scale of relevant
        magnitude. If set to `"latent"`, use the latent libary size.
    n_samples
        Number of posterior samples to use for estimation.
    batch_size
        Minibatch size for data loading into model. Defaults to `scvi.settings.batch_size`.
    return_mean
        Whether to return the mean of the samples.
    return_numpy
        Return a :class:`~numpy.ndarray` instead of a :class:`~pandas.DataFrame`. DataFrame includes
        gene names as columns. If either `n_samples=1` or `return_mean=True`, defaults to `False`.
        Otherwise, it defaults to `True`.

    Returns
    -------
    If `n_samples` > 1 and `return_mean` is False, then the shape is `(samples, cells, genes)`.
    Otherwise, shape is `(cells, genes)`. In this case, return type is :class:`~pandas.DataFrame` unless `return_numpy` is True.
    """
    adata = self._validate_anndata(adata)
    scdl = self._make_scvi_dl(adata=adata, indices=indices, batch_size=batch_size)

    transform_batch = _get_batch_code_from_category(adata, transform_batch)

    if gene_list is None:
        gene_mask = slice(None)
    else:
        all_genes = _get_var_names_from_setup_anndata(adata)
        gene_mask = [True if gene in gene_list else False for gene in all_genes]

    if n_samples > 1 and return_mean is False:
        if return_numpy is False:
            logger.warning(
                "return_numpy must be True if n_samples > 1 and return_mean is False, returning np.ndarray"
            )
        return_numpy = True
    if indices is None:
        indices = np.arange(adata.n_obs)
    if library_size == "latent":
        model_fn = self.model.get_sample_rate
        scaling = 1
    else:
        model_fn = self.model.get_sample_scale
        scaling = library_size

    exprs = []
    for tensors in scdl:
        x = tensors[_CONSTANTS.X_KEY]
        batch_idx = tensors[_CONSTANTS.BATCH_KEY]
        labels = tensors[_CONSTANTS.LABELS_KEY]
        per_batch_exprs = []
        for batch in transform_batch:
            output = model_fn(
                x,
                batch_index=batch_idx,
                y=labels,
                n_samples=n_samples,
                transform_batch=batch,
            )[..., gene_mask]
            output *= scaling
            output = output.cpu().numpy()
            per_batch_exprs.append(output)
        per_batch_exprs = np.stack(
            per_batch_exprs
        )  # shape is (len(transform_batch) x batch_size x n_var)
        exprs += [per_batch_exprs.mean(0)]

    if n_samples > 1:
        # The -2 axis correspond to cells.
        exprs = np.concatenate(exprs, axis=-2)
    else:
        exprs = np.concatenate(exprs, axis=0)
    if n_samples > 1 and return_mean:
        exprs = exprs.mean(0)

    if return_numpy is None or return_numpy is False:
        return pd.DataFrame(
            exprs,
            columns=adata.var_names[gene_mask],
            index=adata.obs_names[indices],
        )
    else:
        return exprs