Repository: scverse/scanpy

Issue:
Broken zappy tests
### Please make sure these conditions are met

- [X] I have checked that this issue has not already been reported.
- [X] I have confirmed this bug exists on the latest version of scanpy.
- [ ] (optional) I have confirmed this bug exists on the main branch of scanpy.

### What happened?

Somewhere, we broke zappy tests.

This diff is the start of fixing this: https://github.com/scverse/scanpy/compare/d3a8c7981321e4f0afa3290dd84c06fed3654b0a..fd90edb71b6f5fc4cc58638583ea4ff37b34e660

### Minimal code sample

```python
pip install zappy
pytest scanpy/tests/test_preprocessing_distributed.py
```


### Error output

```pytb
____________________________________________________________________________________________________________________________________ test_normalize_total[direct] ____________________________________________________________________________________________________________________________________

adata = AnnData object with n_obs × n_vars = 10000 × 1000
    var: 'gene_ids', adata_dist = AnnData object with n_obs × n_vars = 10000 × 1000
    var: 'gene_ids'
    uns: 'dist-mode'

    def test_normalize_total(adata: AnnData, adata_dist: AnnData):
>       normalize_total(adata_dist)

scanpy/tests/test_preprocessing_distributed.py:93: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
/home/phil/.local/share/hatch/env/virtual/scanpy/q4In3tK-/hatch-test.stable/lib/python3.12/site-packages/legacy_api_wrap/__init__.py:80: in fn_compatible
    return fn(*args_all, **kw)
scanpy/preprocessing/_normalization.py:240: in normalize_total
    adata, _normalize_data(X, counts_per_cell, target_sum), layer=layer
scanpy/preprocessing/_normalization.py:49: in _normalize_data
    return axis_mul_or_truediv(
/usr/lib/python3.12/functools.py:909: in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

X = <zappy.direct.array.DirectZappyArray object at 0x7f06eb607a40>, scaling_array = array([0.9270073, 0.4379562, 1.5985402, ..., 0.8540146, 1.5109489,
       0.8978102], dtype=float32), axis = 0, op = <built-in function truediv>

    @singledispatch
    def axis_mul_or_truediv(
        X: sparse.spmatrix,
        scaling_array,
        axis: Literal[0, 1],
        op: Callable[[Any, Any], Any],
        *,
        allow_divide_by_zero: bool = True,
        out: sparse.spmatrix | None = None,
    ) -> sparse.spmatrix:
        check_op(op)
        if out is not None:
            if X.data is not out.data:
                raise ValueError(
                    "`out` argument provided but not equal to X.  This behavior is not supported for sparse matrix scaling."
                )
        if not allow_divide_by_zero and op is truediv:
            scaling_array = scaling_array.copy() + (scaling_array == 0)
    
        row_scale = axis == 0
        column_scale = axis == 1
        if row_scale:
    
            def new_data_op(x):
                return op(x.data, np.repeat(scaling_array, np.diff(x.indptr)))
    
        elif column_scale:
    
            def new_data_op(x):
                return op(x.data, scaling_array.take(x.indices, mode="clip"))
    
>       if X.format == "csr":
E       AttributeError: 'DirectZappyArray' object has no attribute 'format'

scanpy/_utils/__init__.py:612: AttributeError
```


### Versions

<details>

```
-----
anndata     0.10.7
scanpy      1.10.0rc2.dev61+g874d99b3
-----
PIL                 10.3.0
asciitree           NA
cloudpickle         3.0.0
coverage            7.5.3
cycler              0.12.1
cython_runtime      NA
dask                2024.5.2
dateutil            2.9.0.post0
distutils           3.12.3
h5py                3.11.0
igraph              0.11.5
joblib              1.4.2
kiwisolver          1.4.5
legacy_api_wrap     NA
leidenalg           0.10.2
llvmlite            0.42.0
matplotlib          3.9.0
mpl_toolkits        NA
natsort             8.4.0
numba               0.59.1
numcodecs           0.12.1
numpy               1.26.4
packaging           24.0
pandas              2.2.2
psutil              5.9.8
pyparsing           3.1.2
pytz                2024.1
scipy               1.13.1
session_info        1.0.0
setuptools          70.0.0
setuptools_scm      NA
six                 1.16.0
sklearn             1.5.0
texttable           1.7.0
threadpoolctl       3.5.0
tlz                 0.12.1
toolz               0.12.1
yaml                6.0.1
zappy               NA
zarr                2.18.2
-----
Python 3.12.3 (main, Apr 23 2024, 09:16:07) [GCC 13.2.1 20240417]
Linux-6.9.3-zen1-1-zen-x86_64-with-glibc2.39
-----
Session information updated at 2024-06-03 11:27
```

</details>



Discussion:

Resolving patch:
diff --git a/scanpy/_utils/__init__.py b/scanpy/_utils/__init__.py
--- a/scanpy/_utils/__init__.py
+++ b/scanpy/_utils/__init__.py
@@ -580,14 +580,34 @@ def check_op(op):
 
 @singledispatch
 def axis_mul_or_truediv(
-    X: sparse.spmatrix,
+    X: np.ndarray,
+    scaling_array: np.ndarray,
+    axis: Literal[0, 1],
+    op: Callable[[Any, Any], Any],
+    *,
+    allow_divide_by_zero: bool = True,
+    out: np.ndarray | None = None,
+) -> np.ndarray:
+    check_op(op)
+    scaling_array = broadcast_axis(scaling_array, axis)
+    if op is mul:
+        return np.multiply(X, scaling_array, out=out)
+    if not allow_divide_by_zero:
+        scaling_array = scaling_array.copy() + (scaling_array == 0)
+    return np.true_divide(X, scaling_array, out=out)
+
+
+@axis_mul_or_truediv.register(sparse.csr_matrix)
+@axis_mul_or_truediv.register(sparse.csc_matrix)
+def _(
+    X: sparse.csr_matrix | sparse.csc_matrix,
     scaling_array,
     axis: Literal[0, 1],
     op: Callable[[Any, Any], Any],
     *,
     allow_divide_by_zero: bool = True,
-    out: sparse.spmatrix | None = None,
-) -> sparse.spmatrix:
+    out: sparse.csr_matrix | sparse.csc_matrix | None = None,
+) -> sparse.csr_matrix | sparse.csc_matrix:
     check_op(op)
     if out is not None:
         if X.data is not out.data:
@@ -629,25 +649,6 @@ def new_data_op(x):
     ).T
 
 
-@axis_mul_or_truediv.register(np.ndarray)
-def _(
-    X: np.ndarray,
-    scaling_array: np.ndarray,
-    axis: Literal[0, 1],
-    op: Callable[[Any, Any], Any],
-    *,
-    allow_divide_by_zero: bool = True,
-    out: np.ndarray | None = None,
-) -> np.ndarray:
-    check_op(op)
-    scaling_array = broadcast_axis(scaling_array, axis)
-    if op is mul:
-        return np.multiply(X, scaling_array, out=out)
-    if not allow_divide_by_zero:
-        scaling_array = scaling_array.copy() + (scaling_array == 0)
-    return np.true_divide(X, scaling_array, out=out)
-
-
 def make_axis_chunks(
     X: DaskArray, axis: Literal[0, 1], pad=True
 ) -> tuple[tuple[int], tuple[int]]:


Test patch: