# Buggy code
def _check_anndata_setup_equivalence(adata_source, adata_target):
    """Checks if target setup is equivalent to source."""
    if isinstance(adata_source, anndata.AnnData):
        _scvi_dict = adata_source.uns["_scvi"]
    else:
        _scvi_dict = adata_source
    adata = adata_target

    stats = _scvi_dict["summary_stats"]
    use_raw = _scvi_dict["use_raw"]

    target_n_vars = adata.shape[1] if not use_raw else adata.raw.shape[1]
    error_msg = (
        "Number of {} in anndata different from initial anndata used for training."
    )
    if target_n_vars != stats["n_vars"]:
        raise ValueError(error_msg.format("vars"))

    error_msg = (
        "There are more {} categories in the data than were originally registered. "
        + "Please check your {} categories as well as adata.uns['_scvi']['categorical_mappings']."
    )
    self_categoricals = _scvi_dict["categorical_mappings"]
    self_batch_mapping = self_categoricals["_scvi_batch"]["mapping"]

    adata_categoricals = adata.uns["_scvi"]["categorical_mappings"]
    adata_batch_mapping = adata_categoricals["_scvi_batch"]["mapping"]
    # check if the categories are the same
    error_msg = (
        "Categorial encoding for {} is not the same between "
        + "the anndata used to train the model and the anndata just passed in. "
        + "Categorical encoding needs to be same elements, same order, and same datatype.\n"
        + "Expected categories: {}. Received categories: {}.\n"
        + "Try running `dataset.transfer_anndata_setup()` or deleting `adata.uns['_scvi']."
    )
    if not _assert_equal_mapping(self_batch_mapping, adata_batch_mapping):
        raise ValueError(
            error_msg.format("batch", self_batch_mapping, adata_batch_mapping)
        )
    self_labels_mapping = self_categoricals["_scvi_labels"]["mapping"]
    adata_labels_mapping = adata_categoricals["_scvi_labels"]["mapping"]
    if not _assert_equal_mapping(self_labels_mapping, adata_labels_mapping):
        raise ValueError(
            error_msg.format("label", self_labels_mapping, adata_labels_mapping)
        )

    # validate any extra categoricals
    if "extra_categorical_mappings" in _scvi_dict.keys():
        target_extra_cat_maps = adata.uns["_scvi"]["extra_categorical_mappings"]
        for key, val in _scvi_dict["extra_categorical_mappings"].items():
            target_map = target_extra_cat_maps[key]
            if not _assert_equal_mapping(val, target_map):
                raise ValueError(error_msg.format(key, val, target_map))
    # validate any extra continuous covs
    if "extra_continuous_keys" in _scvi_dict.keys():
        if "extra_continuous_keys" not in adata.uns["_scvi"].keys():
            raise ValueError('extra_continuous_keys not in adata.uns["_scvi"]')
        target_cont_keys = adata.uns["_scvi"]["extra_continuous_keys"]
        if not _scvi_dict["extra_continuous_keys"].equals(target_cont_keys):
            raise ValueError(
                "extra_continous_keys are not the same between source and target"
            )


# 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 _check_anndata_setup_equivalence(adata_source, adata_target):
    """Checks if target setup is equivalent to source."""
    if isinstance(adata_source, anndata.AnnData):
        _scvi_dict = adata_source.uns["_scvi"]
    else:
        _scvi_dict = adata_source
    adata = adata_target

    stats = _scvi_dict["summary_stats"]

    target_n_vars = adata.shape[1]
    error_msg = (
        "Number of {} in anndata different from initial anndata used for training."
    )
    if target_n_vars != stats["n_vars"]:
        raise ValueError(error_msg.format("vars"))

    error_msg = (
        "There are more {} categories in the data than were originally registered. "
        + "Please check your {} categories as well as adata.uns['_scvi']['categorical_mappings']."
    )
    self_categoricals = _scvi_dict["categorical_mappings"]
    self_batch_mapping = self_categoricals["_scvi_batch"]["mapping"]

    adata_categoricals = adata.uns["_scvi"]["categorical_mappings"]
    adata_batch_mapping = adata_categoricals["_scvi_batch"]["mapping"]
    # check if the categories are the same
    error_msg = (
        "Categorial encoding for {} is not the same between "
        + "the anndata used to train the model and the anndata just passed in. "
        + "Categorical encoding needs to be same elements, same order, and same datatype.\n"
        + "Expected categories: {}. Received categories: {}.\n"
        + "Try running `dataset.transfer_anndata_setup()` or deleting `adata.uns['_scvi']."
    )
    if not _assert_equal_mapping(self_batch_mapping, adata_batch_mapping):
        raise ValueError(
            error_msg.format("batch", self_batch_mapping, adata_batch_mapping)
        )
    self_labels_mapping = self_categoricals["_scvi_labels"]["mapping"]
    adata_labels_mapping = adata_categoricals["_scvi_labels"]["mapping"]
    if not _assert_equal_mapping(self_labels_mapping, adata_labels_mapping):
        raise ValueError(
            error_msg.format("label", self_labels_mapping, adata_labels_mapping)
        )

    # validate any extra categoricals
    if "extra_categorical_mappings" in _scvi_dict.keys():
        target_extra_cat_maps = adata.uns["_scvi"]["extra_categorical_mappings"]
        for key, val in _scvi_dict["extra_categorical_mappings"].items():
            target_map = target_extra_cat_maps[key]
            if not _assert_equal_mapping(val, target_map):
                raise ValueError(error_msg.format(key, val, target_map))
    # validate any extra continuous covs
    if "extra_continuous_keys" in _scvi_dict.keys():
        if "extra_continuous_keys" not in adata.uns["_scvi"].keys():
            raise ValueError('extra_continuous_keys not in adata.uns["_scvi"]')
        target_cont_keys = adata.uns["_scvi"]["extra_continuous_keys"]
        if not _scvi_dict["extra_continuous_keys"].equals(target_cont_keys):
            raise ValueError(
                "extra_continous_keys are not the same between source and target"
            )