# Buggy code
def _compute_library_size_batch(
    adata,
    batch_key: str,
    local_l_mean_key: str = None,
    local_l_var_key: str = None,
    layer=None,
    use_raw=False,
    copy: bool = False,
):
    """
    Computes the library size.

    Parameters
    ----------
    adata
        anndata object containing counts
    batch_key
        key in obs for batch information
    local_l_mean_key
        key in obs to save the local log mean
    local_l_var_key
        key in obs to save the local log variance
    layer
        if not None, will use this in adata.layers[] for X
    use_raw
        Use ``.raw`` for X
    copy
        if True, returns a copy of the adata

    Returns
    -------
    type
        anndata.AnnData if copy was True, else None

    """
    if batch_key not in adata.obs_keys():
        raise ValueError("batch_key not valid key in obs dataframe")
    local_means = np.zeros((adata.shape[0], 1))
    local_vars = np.zeros((adata.shape[0], 1))
    batch_indices = adata.obs[batch_key]
    for i_batch in np.unique(batch_indices):
        idx_batch = np.squeeze(batch_indices == i_batch)
        if use_raw:
            data = adata[idx_batch].raw.X
        elif layer is not None:
            if layer not in adata.layers.keys():
                raise ValueError("layer not a valid key for adata.layers")
            data = adata[idx_batch].layers[layer]
        else:
            data = adata[idx_batch].X
        (local_means[idx_batch], local_vars[idx_batch]) = _compute_library_size(data)
    if local_l_mean_key is None:
        local_l_mean_key = "_scvi_local_l_mean"
    if local_l_var_key is None:
        local_l_var_key = "_scvi_local_l_var"

    if copy:
        copy = adata.copy()
        copy.obs[local_l_mean_key] = local_means
        copy.obs[local_l_var_key] = local_vars
        return copy
    else:
        adata.obs[local_l_mean_key] = local_means
        adata.obs[local_l_var_key] = local_vars


# Traceback (AttributeError)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-125-eb1e0fcfe1d1> in <module>
1 de = model.differential_expression(
----> 2     groupby="leiden",
3 )
4 de.head()

~/anaconda3/lib/python3.7/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)
179         adata = self._validate_anndata(adata)
180
--> 181         col_names = _get_var_names_from_setup_anndata(adata)
182         model_fn = partial(
183             self.get_normalized_expression,

~/anaconda3/lib/python3.7/site-packages/scvi/model/_utils.py in _get_var_names_from_setup_anndata(adata)
119         adata.var_names
120         if adata.uns["_scvi"]["use_raw"] is False
--> 121         else adata.raw.var_names
122     )
123

AttributeError: 'NoneType' object has no attribute 'var_names'

# Fixed code
def _compute_library_size_batch(
    adata,
    batch_key: str,
    local_l_mean_key: str = None,
    local_l_var_key: str = None,
    layer=None,
    copy: bool = False,
):
    """
    Computes the library size.

    Parameters
    ----------
    adata
        anndata object containing counts
    batch_key
        key in obs for batch information
    local_l_mean_key
        key in obs to save the local log mean
    local_l_var_key
        key in obs to save the local log variance
    layer
        if not None, will use this in adata.layers[] for X
    copy
        if True, returns a copy of the adata

    Returns
    -------
    type
        anndata.AnnData if copy was True, else None

    """
    if batch_key not in adata.obs_keys():
        raise ValueError("batch_key not valid key in obs dataframe")
    local_means = np.zeros((adata.shape[0], 1))
    local_vars = np.zeros((adata.shape[0], 1))
    batch_indices = adata.obs[batch_key]
    for i_batch in np.unique(batch_indices):
        idx_batch = np.squeeze(batch_indices == i_batch)
        if layer is not None:
            if layer not in adata.layers.keys():
                raise ValueError("layer not a valid key for adata.layers")
            data = adata[idx_batch].layers[layer]
        else:
            data = adata[idx_batch].X
        (local_means[idx_batch], local_vars[idx_batch]) = _compute_library_size(data)
    if local_l_mean_key is None:
        local_l_mean_key = "_scvi_local_l_mean"
    if local_l_var_key is None:
        local_l_var_key = "_scvi_local_l_var"

    if copy:
        copy = adata.copy()
        copy.obs[local_l_mean_key] = local_means
        copy.obs[local_l_var_key] = local_vars
        return copy
    else:
        adata.obs[local_l_mean_key] = local_means
        adata.obs[local_l_var_key] = local_vars