feat: add Array API support via array-api-compat#4179
Conversation
|
flying-sheep
left a comment
flying-sheep
left a comment
There was a problem hiding this comment.
OK, please remember the comments from the fast-array-utils PR, e.g. we agreed to handle singledispatch functions like this: https://github.com/scverse/fast-array-utils/blob/febaf245ecd6da849c2331638219f649935d5a21/src/fast_array_utils/stats/_power.py#L38-L39
We also need tests. There’s MAP_ARRAY_TYPES which should gain an entry, which will make many many tests run with the type you add. We can discuss how to best exclude certain tests, for starters you can just add it and make some of the tests fail – that’ll show you where things aren’t working yet (and might or might not be fixed in this PR)
Please also add a release note (hatch run towncrier:create 4179.feat.md)
for more information, see https://pre-commit.ci
…lia-k510/scanpy into array-api-compat-integration
for more information, see https://pre-commit.ci
flying-sheep
changed the title
| @axis_nnz.register(HasArrayNamespace) | ||
| def _(x: HasArrayNamespace, /, axis: Literal[0, 1]) -> Any: | ||
|
|
||
| xp = get_namespace(x) | ||
| return xp.count_nonzero(x, axis=axis) |
There was a problem hiding this comment.
this handles numpy with equal performace too no?
| if pkg_version("scipy") >= Version("1.15"): | ||
| # newer scipy versions support the `axis` argument for count_nonzero | ||
| @axis_nnz.register(CSBase) | ||
| def _(x: CSBase, /, axis: Literal[0, 1]) -> np.ndarray: | ||
| return x.count_nonzero(axis=axis) | ||
|
|
||
| else: | ||
| # older scipy versions don’t have any way to get the nnz of a sparse array | ||
| @axis_nnz.register(CSBase) | ||
| def _(x: CSBase, /, axis: Literal[0, 1]) -> np.ndarray: | ||
| if isinstance(x, _CSArray): | ||
| from scipy.sparse import csc_array, csr_array # noqa: TID251 | ||
|
|
||
| x = (csr_array if x.format == "csr" else csc_array)(x) | ||
| return x.getnnz(axis=axis) | ||
|
|
There was a problem hiding this comment.
looks like you got a merge wrong, this change reverts an intentional simplification.
| @check_nonnegative_integers.register(HasArrayNamespace) | ||
| def _check_nonnegative_integers_array_api(x: HasArrayNamespace, /) -> bool: |
There was a problem hiding this comment.
should this go below the np.ndarray case? I think _check_nonnegative_integers_array_api.dispatch(np.ndarray) will return the first match, no?
| @_resolve_vals.register(HasArrayNamespace) | ||
| def _resolve_vals_array_api(val: HasArrayNamespace) -> NDArray: | ||
| # Moran's I / Geary's C use numba kernels, so convert at the boundary | ||
| return np.asarray(val) |
There was a problem hiding this comment.
same here, should go below the np.ndarray case right?
| if n_removed: | ||
| x = x[:, filt].copy() | ||
|
|
||
| from .._compat import get_namespace ### double check |
There was a problem hiding this comment.
what does ### double check mean?
| counts = counts.compute() | ||
|
|
||
| counts_greater_than_zero = counts[counts > 0] | ||
| median = np.median(counts_greater_than_zero) |
| logg.info( | ||
| "... as scaling leads to float results, integer " | ||
| "input is cast to float, returning copy." | ||
| ) | ||
| x = xp.astype(x, xp.float64) |
There was a problem hiding this comment.
assign the message before this if/else block so it’s not duplicated
| return log1p_array(data, copy=copy, base=base) | ||
|
|
||
|
|
||
| @log1p.register(HasArrayNamespace) |
There was a problem hiding this comment.
again specificity/order. please check all of these in case I forgot one!
| return metadata(package) | ||
|
|
||
|
|
||
| def get_namespace(x): |
There was a problem hiding this comment.
please add the appropriate return type
| if find_spec("jax"): | ||
| import jax | ||
|
|
||
| jax.config.update("jax_enable_x64", True) # noqa: FBT003 |
There was a problem hiding this comment.
can you add a short comment what that does?
2 participants
This adds Array API support to scanpy's preprocessing pipeline, where I am planning to use JAX as the test case. The goal is for someone to put a JAX array into
adata.Xand run the usual pipeline without the array getting silently pulled to CPU partway through. It builds on the recentfast-array-utilswork, which means a lot of the heavy lifting (sums, means, variances) already works across backends.