# Buggy code
def get_feature_correlation_matrix(
    self,
    adata: Optional[AnnData] = None,
    indices: Optional[Sequence[int]] = None,
    n_samples: int = 10,
    batch_size: int = 64,
    rna_size_factor: int = 1000,
    transform_batch: Optional[Sequence[Union[Number, str]]] = None,
    correlation_type: Literal["spearman", "pearson"] = "spearman",
) -> pd.DataFrame:
    """
    Generate gene-gene correlation matrix using scvi uncertainty and expression.

    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.
    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`.
    rna_size_factor
        size factor for RNA prior to sampling gamma distribution.
    transform_batch
        Batches to condition on.
        If transform_batch is:

        - None, then real observed batch is used.
        - int, then batch transform_batch is used.
        - list of int, then values are averaged over provided batches.
    correlation_type
        One of "pearson", "spearman".

    Returns
    -------
    Gene-gene correlation matrix
    """
    from scipy.stats import spearmanr

    adata = self._validate_anndata(adata)

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

    transform_batch = _get_batch_code_from_category(adata, transform_batch)

    corr_mats = []
    for b in transform_batch:
        denoised_data = self._get_denoised_samples(
            adata=adata,
            indices=indices,
            n_samples=n_samples,
            batch_size=batch_size,
            rna_size_factor=rna_size_factor,
            transform_batch=b,
        )
        flattened = np.zeros(
            (denoised_data.shape[0] * n_samples, denoised_data.shape[1])
        )
        for i in range(n_samples):
            flattened[
                denoised_data.shape[0] * (i) : denoised_data.shape[0] * (i + 1)
            ] = denoised_data[:, :, i]
        if correlation_type == "pearson":
            corr_matrix = np.corrcoef(flattened, rowvar=False)
        elif correlation_type == "spearman":
            corr_matrix, _ = spearmanr(flattened)
        else:
            raise ValueError(
                "Unknown correlation type. Choose one of 'spearman', 'pearson'."
            )
        corr_mats.append(corr_matrix)
    corr_matrix = np.mean(np.stack(corr_mats), axis=0)
    var_names = _get_var_names_from_setup_anndata(adata)
    return pd.DataFrame(corr_matrix, index=var_names, columns=var_names)


# 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_feature_correlation_matrix(
    self,
    adata: Optional[AnnData] = None,
    indices: Optional[Sequence[int]] = None,
    n_samples: int = 10,
    batch_size: int = 64,
    rna_size_factor: int = 1000,
    transform_batch: Optional[Sequence[Union[Number, str]]] = None,
    correlation_type: Literal["spearman", "pearson"] = "spearman",
) -> pd.DataFrame:
    """
    Generate gene-gene correlation matrix using scvi uncertainty and expression.

    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.
    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`.
    rna_size_factor
        size factor for RNA prior to sampling gamma distribution.
    transform_batch
        Batches to condition on.
        If transform_batch is:

        - None, then real observed batch is used.
        - int, then batch transform_batch is used.
        - list of int, then values are averaged over provided batches.
    correlation_type
        One of "pearson", "spearman".

    Returns
    -------
    Gene-gene correlation matrix
    """
    from scipy.stats import spearmanr

    adata = self._validate_anndata(adata)

    transform_batch = _get_batch_code_from_category(adata, transform_batch)

    corr_mats = []
    for b in transform_batch:
        denoised_data = self._get_denoised_samples(
            adata=adata,
            indices=indices,
            n_samples=n_samples,
            batch_size=batch_size,
            rna_size_factor=rna_size_factor,
            transform_batch=b,
        )
        flattened = np.zeros(
            (denoised_data.shape[0] * n_samples, denoised_data.shape[1])
        )
        for i in range(n_samples):
            flattened[
                denoised_data.shape[0] * (i) : denoised_data.shape[0] * (i + 1)
            ] = denoised_data[:, :, i]
        if correlation_type == "pearson":
            corr_matrix = np.corrcoef(flattened, rowvar=False)
        elif correlation_type == "spearman":
            corr_matrix, _ = spearmanr(flattened)
        else:
            raise ValueError(
                "Unknown correlation type. Choose one of 'spearman', 'pearson'."
            )
        corr_mats.append(corr_matrix)
    corr_matrix = np.mean(np.stack(corr_mats), axis=0)
    var_names = _get_var_names_from_setup_anndata(adata)
    return pd.DataFrame(corr_matrix, index=var_names, columns=var_names)