# Buggy code
def scale_sampler(
    self,
    selection: Union[List[bool], np.ndarray],
    n_samples: Optional[int] = 5000,
    n_samples_per_cell: Optional[int] = None,
    batchid: Optional[Union[List[int], np.ndarray]] = None,
    use_observed_batches: Optional[bool] = False,
    give_mean: Optional[bool] = False,
) -> dict:
    """
    Samples the posterior scale using the variational posterior distribution.

    Parameters
    ----------
    selection
        Mask or list of cell ids to select
    n_samples
        Number of samples in total per batch (fill either `n_samples_total`
        or `n_samples_per_cell`)
    n_samples_per_cell
        Number of time we sample from each observation per batch
        (fill either `n_samples_total` or `n_samples_per_cell`)
    batchid
        Biological batch for which to sample from.
        Default (None) sample from all batches
    use_observed_batches
        Whether normalized means are conditioned on observed
        batches or if observed batches are to be used
    give_mean
        Return mean of values


    Returns
    -------
    type
        Dictionary containing:
        `scale`
        Posterior aggregated scale samples of shape (n_samples, n_vars)
        where n_samples correspond to either:
        - n_bio_batches * n_cells * n_samples_per_cell
        or
        - n_samples_total
        `batch`
        associated batch ids

    """
    # Get overall number of desired samples and desired batches
    if batchid is None and not use_observed_batches:
        # TODO determine if we iterate over all categorical batches from train dataset
        # or just the batches in adata
        batchid = np.unique(get_from_registry(self.adata, key=_CONSTANTS.BATCH_KEY))
    if use_observed_batches:
        if batchid is not None:
            raise ValueError("Unconsistent batch policy")
        batchid = [None]
    if n_samples is None and n_samples_per_cell is None:
        n_samples = 5000
    elif n_samples_per_cell is not None and n_samples is None:
        n_samples = n_samples_per_cell * len(selection)
    if (n_samples_per_cell is not None) and (n_samples is not None):
        warnings.warn(
            "n_samples and n_samples_per_cell were provided. Ignoring n_samples_per_cell"
        )
    n_samples = int(n_samples / len(batchid))
    if n_samples == 0:
        warnings.warn("very small sample size, please consider increasing `n_samples`")
        n_samples = 2

    # Selection of desired cells for sampling
    if selection is None:
        raise ValueError("selections should be a list of cell subsets indices")
    selection = np.asarray(selection)
    if selection.dtype is np.dtype("bool"):
        if len(selection) < self.adata.shape[0]:
            raise ValueError("Mask must be same length as adata.")
        selection = np.asarray(np.where(selection)[0].ravel())

    # Sampling loop
    px_scales = []
    batch_ids = []
    for batch_idx in batchid:
        idx = np.random.choice(np.arange(self.adata.shape[0])[selection], n_samples)
        px_scales.append(
            self.model_fn(self.adata, indices=idx, transform_batch=batch_idx)
        )
        batch_idx = batch_idx if batch_idx is not None else np.nan
        batch_ids.append([batch_idx] * px_scales[-1].shape[0])
    px_scales = np.concatenate(px_scales)
    batch_ids = np.concatenate(batch_ids).reshape(-1)
    if px_scales.shape[0] != batch_ids.shape[0]:
        raise ValueError("sampled scales and batches have inconsistent shapes")
    if give_mean:
        px_scales = px_scales.mean(0)
    return dict(scale=px_scales, batch=batch_ids)


# 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 scale_sampler(
    self,
    selection: Union[List[bool], np.ndarray],
    n_samples: Optional[int] = 5000,
    n_samples_per_cell: Optional[int] = None,
    batchid: Optional[Sequence[Union[Number, str]]] = None,
    use_observed_batches: Optional[bool] = False,
    give_mean: Optional[bool] = False,
) -> dict:
    """
    Samples the posterior scale using the variational posterior distribution.

    Parameters
    ----------
    selection
        Mask or list of cell ids to select
    n_samples
        Number of samples in total per batch (fill either `n_samples_total`
        or `n_samples_per_cell`)
    n_samples_per_cell
        Number of time we sample from each observation per batch
        (fill either `n_samples_total` or `n_samples_per_cell`)
    batchid
        Biological batch for which to sample from.
        Default (None) sample from all batches
    use_observed_batches
        Whether normalized means are conditioned on observed
        batches or if observed batches are to be used
    give_mean
        Return mean of values


    Returns
    -------
    type
        Dictionary containing:
        `scale`
        Posterior aggregated scale samples of shape (n_samples, n_vars)
        where n_samples correspond to either:
        - n_bio_batches * n_cells * n_samples_per_cell
        or
        - n_samples_total
        `batch`
        associated batch ids

    """
    # Get overall number of desired samples and desired batches
    if batchid is None and not use_observed_batches:
        categorical_mappings = self.adata.uns["_scvi"]["categorical_mappings"]
        batchid = categorical_mappings["_scvi_batch"]["mapping"]
    if use_observed_batches:
        if batchid is not None:
            raise ValueError("Unconsistent batch policy")
        batchid = [None]
    if n_samples is None and n_samples_per_cell is None:
        n_samples = 5000
    elif n_samples_per_cell is not None and n_samples is None:
        n_samples = n_samples_per_cell * len(selection)
    if (n_samples_per_cell is not None) and (n_samples is not None):
        warnings.warn(
            "n_samples and n_samples_per_cell were provided. Ignoring n_samples_per_cell"
        )
    n_samples = int(n_samples / len(batchid))
    if n_samples == 0:
        warnings.warn("very small sample size, please consider increasing `n_samples`")
        n_samples = 2

    # Selection of desired cells for sampling
    if selection is None:
        raise ValueError("selections should be a list of cell subsets indices")
    selection = np.asarray(selection)
    if selection.dtype is np.dtype("bool"):
        if len(selection) < self.adata.shape[0]:
            raise ValueError("Mask must be same length as adata.")
        selection = np.asarray(np.where(selection)[0].ravel())

    # Sampling loop
    px_scales = []
    batch_ids = []
    for batch_idx in batchid:
        idx = np.random.choice(np.arange(self.adata.shape[0])[selection], n_samples)
        px_scales.append(
            self.model_fn(self.adata, indices=idx, transform_batch=batch_idx)
        )
        batch_idx = batch_idx if batch_idx is not None else np.nan
        batch_ids.append([batch_idx] * px_scales[-1].shape[0])
    px_scales = np.concatenate(px_scales)
    batch_ids = np.concatenate(batch_ids).reshape(-1)
    if px_scales.shape[0] != batch_ids.shape[0]:
        raise ValueError("sampled scales and batches have inconsistent shapes")
    if give_mean:
        px_scales = px_scales.mean(0)
    return dict(scale=px_scales, batch=batch_ids)