# Buggy code
def _de_core(
    adata,
    model_fn,
    groupby,
    group1,
    group2,
    idx1,
    idx2,
    all_stats,
    all_stats_fn,
    col_names,
    mode,
    batchid1,
    batchid2,
    delta,
    batch_correction,
    **kwargs,
):
    """Internal function for DE interface."""
    if group1 is None and idx1 is None:
        group1 = adata.obs[groupby].cat.categories.tolist()

    if isinstance(group1, str):
        group1 = [group1]

    # make a temp obs key using indices
    temp_key = None
    if idx1 is not None:
        idx1 = np.asarray(idx1).ravel()
        g1_key = "one"
        obs_col = np.array(["None"] * adata.shape[0], dtype=str)
        obs_col[idx1] = g1_key
        group2 = None if idx2 is None else "two"
        if idx2 is not None:
            idx2 = np.asarray(idx2).ravel()
            obs_col[idx2] = group2
        temp_key = "_scvi_temp_de"
        adata.obs[temp_key] = obs_col
        groupby = temp_key
        group1 = [g1_key]

    df_results = []
    dc = DifferentialComputation(model_fn, adata)
    for g1 in track(
        group1,
        description="DE...",
    ):
        cell_idx1 = (adata.obs[groupby] == g1).ravel()
        if group2 is None:
            cell_idx2 = ~cell_idx1
        else:
            cell_idx2 = adata.obs[groupby] == group2

        all_info = dc.get_bayes_factors(
            cell_idx1,
            cell_idx2,
            mode=mode,
            delta=delta,
            batchid1=batchid1,
            batchid2=batchid2,
            use_observed_batches=not batch_correction,
            **kwargs,
        )

        if all_stats is True:
            genes_properties_dict = all_stats_fn(adata, cell_idx1, cell_idx2)
            all_info = {**all_info, **genes_properties_dict}

        res = pd.DataFrame(all_info, index=col_names)
        sort_key = "proba_de" if mode == "change" else "bayes_factor"
        res = res.sort_values(by=sort_key, ascending=False)
        if idx1 is None:
            g2 = "Rest" if group2 is None else group2
            res["comparison"] = "{} vs {}".format(g1, g2)
        df_results.append(res)

    if temp_key is not None:
        del adata.obs[temp_key]

    result = pd.concat(df_results, axis=0)

    return result


# Traceback (ValueError)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-17-d7b37e40fc13> in <module>()
5     idx1=idx1,
6     idx2=idx2,
----> 7     mode='change',
8 )

7 frames
/usr/local/lib/python3.6/dist-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)
202             delta,
203             batch_correction,
--> 204             **kwargs,
205         )
206

/usr/local/lib/python3.6/dist-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)
70
71         if all_stats is True:
---> 72             genes_properties_dict = all_stats_fn(adata, cell_idx1, cell_idx2)
73             all_info = {**all_info, **genes_properties_dict}
74

/usr/local/lib/python3.6/dist-packages/scvi/model/_utils.py in scrna_raw_counts_properties(adata, idx1, idx2)
35     data = get_from_registry(adata, _CONSTANTS.X_KEY)
36     data1 = data[idx1]
---> 37     data2 = data[idx2]
38     mean1 = np.asarray((data1).mean(axis=0)).ravel()
39     mean2 = np.asarray((data2).mean(axis=0)).ravel()

/usr/local/lib/python3.6/dist-packages/scipy/sparse/_index.py in __getitem__(self, key)
57                 return self._get_arrayXint(row, col)
58             elif isinstance(col, slice):
---> 59                 return self._get_arrayXslice(row, col)
60         else:  # row.ndim == 2
61             if isinstance(col, INT_TYPES):

/usr/local/lib/python3.6/dist-packages/scipy/sparse/csr.py in _get_arrayXslice(self, row, col)
323             col = np.arange(*col.indices(self.shape[1]))
324             return self._get_arrayXarray(row, col)
--> 325         return self._major_index_fancy(row)._get_submatrix(minor=col)
326
327

/usr/local/lib/python3.6/dist-packages/scipy/sparse/compressed.py in _major_index_fancy(self, idx)
688         idx_dtype = self.indices.dtype
689         res_indptr = np.zeros(M+1, dtype=idx_dtype)
--> 690         np.cumsum(row_nnz[idx], out=res_indptr[1:])
691
692         nnz = res_indptr[-1]

<__array_function__ internals> in cumsum(*args, **kwargs)

/usr/local/lib/python3.6/dist-packages/numpy/core/fromnumeric.py in cumsum(a, axis, dtype, out)
2468
2469     """
-> 2470     return _wrapfunc(a, 'cumsum', axis=axis, dtype=dtype, out=out)
2471
2472

/usr/local/lib/python3.6/dist-packages/numpy/core/fromnumeric.py in _wrapfunc(obj, method, *args, **kwds)
59
60     try:
---> 61         return bound(*args, **kwds)
62     except TypeError:
63         # A TypeError occurs if the object does have such a method in its

ValueError: provided out is the wrong size for the reduction

# Fixed code
def _de_core(
    adata,
    model_fn,
    groupby,
    group1,
    group2,
    idx1,
    idx2,
    all_stats,
    all_stats_fn,
    col_names,
    mode,
    batchid1,
    batchid2,
    delta,
    batch_correction,
    **kwargs,
):
    """Internal function for DE interface."""
    if group1 is None and idx1 is None:
        group1 = adata.obs[groupby].cat.categories.tolist()

    if isinstance(group1, str):
        group1 = [group1]

    # make a temp obs key using indices
    temp_key = None
    if idx1 is not None:
        idx1 = np.asarray(idx1).ravel()
        g1_key = "one"
        obs_col = np.array(["None"] * adata.shape[0], dtype=str)
        obs_col[idx1] = g1_key
        group2 = None if idx2 is None else "two"
        if idx2 is not None:
            idx2 = np.asarray(idx2).ravel()
            obs_col[idx2] = group2
        temp_key = "_scvi_temp_de"
        adata.obs[temp_key] = obs_col
        groupby = temp_key
        group1 = [g1_key]

    df_results = []
    dc = DifferentialComputation(model_fn, adata)
    for g1 in track(
        group1,
        description="DE...",
    ):
        cell_idx1 = (adata.obs[groupby] == g1).to_numpy().ravel()
        if group2 is None:
            cell_idx2 = ~cell_idx1
        else:
            cell_idx2 = (adata.obs[groupby] == group2).to_numpy().ravel()

        all_info = dc.get_bayes_factors(
            cell_idx1,
            cell_idx2,
            mode=mode,
            delta=delta,
            batchid1=batchid1,
            batchid2=batchid2,
            use_observed_batches=not batch_correction,
            **kwargs,
        )

        if all_stats is True:
            genes_properties_dict = all_stats_fn(adata, cell_idx1, cell_idx2)
            all_info = {**all_info, **genes_properties_dict}

        res = pd.DataFrame(all_info, index=col_names)
        sort_key = "proba_de" if mode == "change" else "bayes_factor"
        res = res.sort_values(by=sort_key, ascending=False)
        if idx1 is None:
            g2 = "Rest" if group2 is None else group2
            res["comparison"] = "{} vs {}".format(g1, g2)
        df_results.append(res)

    if temp_key is not None:
        del adata.obs[temp_key]

    result = pd.concat(df_results, axis=0)

    return result