Adding array-api-compat fallback by amalia-k510 · Pull Request #159 · scverse/fast-array-utils

amalia-k510 · 2026-03-23T10:36:48Z

This PR adds array-api-compat as a fallback in the singledispatch functions across the stats and conv modules so that Array API-compatible arrays (JAX, PyTorch, and others) work out of the box without needing to register each backend individually. The approach is: in each fallback, try array_api_compat.array_namespace(x) first. If it recognizes the array, dispatch through the standard Array API; if not, fall through to the existing numpy path. This touches _generic_ops.py (sum/min/max), _is_constant.py, _power.py, and conv/_to_dense.py. array-api-compat is added as a dependency in pyproject.toml.

Tests are in tests/test_jax.py covering all the stats functions and to_dense with JAX arrays. One thing to flag: _mean_var.py didn't need changes since it calls mean() and power() internally which already go through the fixed dispatchers, at least in my understanding, but JAX requires jax_enable_x64 for the dtype=np.float64 calls to work.

for more information, see https://pre-commit.ci

codecov · 2026-03-23T10:38:16Z

Codecov Report

❌ Patch coverage is 95.65217% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 97.32%. Comparing base (3fac19a) to head (3d4ee3a).

Files with missing lines	Patch %	Lines
src/fast_array_utils/stats/_generic_ops.py	92.30%	1 Missing ⚠️
src/fast_array_utils/stats/_power.py	88.88%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #159      +/-   ##
==========================================
- Coverage   99.29%   97.32%   -1.97%     
==========================================
  Files          21       21              
  Lines         566      598      +32     
==========================================
+ Hits          562      582      +20     
- Misses          4       16      +12

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

codspeed-hq · 2026-03-23T10:46:56Z

Merging this PR will degrade performance by 13.15%

⚠️

Different runtime environments detected

Some benchmarks with significant performance changes were compared across different runtime environments,
which may affect the accuracy of the results.

Open the report in CodSpeed to investigate

❌ 1 regressed benchmark
✅ 231 untouched benchmarks

⚠️ Please fix the performance issues or acknowledge them on CodSpeed.

Performance Changes

	Benchmark	`BASE`	`HEAD`	Efficiency
❌	`test_stats_benchmark[scipy.sparse.csc_array-2d-ax0-float32-is_constant]`	2.9 ms	3.3 ms	-13.15%

_{Comparing amalia-k510:array-api-implementation (3d4ee3a) with main (3fac19a)}

flying-sheep

good start. I wonder if adding a ArrayAPIObject protocol that checks for __array_namespace__ or so could be used instead of putting the array api stuff in the @singledispatch fallback body. What did we do in the AnnData PR?

flying-sheep · 2026-03-23T16:04:02Z

+    if not isinstance(x, np.ndarray) and array_api_compat.is_array_api_obj(x):
+        if to_cpu_memory:
+            return np.asarray(x, order=order)
+        return x  # already dense


I don’t think that follows, but I also don’t know if we can do better.

From what I read, the array API standard only covers dense arrays right now, and sparse types are already caught by their own registered handlers before reaching this fallback. So, from my understanding, if anything that lands here and passes is_array_api_obj is going to be dense. Pydata/Sparse are under active development and once they add sparse support later, we can just register a new handler for it.

It does? I didn’t know it says anything about dense vs sparse. I assumed it just didn’t have any sparse-specific methods defined.

You were right to flag that. I assumed that since the standard focuses on dense arrays, anything reaching this fallback would be dense. But it doesn't actually distinguish dense from sparse, plus, after review, pydata/sparse does implement __array_namespace__, so it would hit this fallback and skip the dense step. I don't think there's a generic way to detect sparsity in the standard either. Should we just register handlers for specific sparse types as they come up, or is there a better way of handling it?

Should we just register handlers for specific sparse types as they come up

yeah I think that’s best for now.

for more information, see https://pre-commit.ci

…k510/fast-array-utils into array-api-implementation

flying-sheep

OK, before we move on, I need to understand this comment:

Catch array-api-compat-wrapped types that lack __array_namespace__ (i.e. PyTorch)

Once I do, I can form an actual opinion about how I’d like this to look (all my comments about an ABC below rely on this). I commented on the line of the comment below so we keep that in a subthread.

Co-authored-by: Philipp A. <flying-sheep@web.de>

for more information, see https://pre-commit.ci

flying-sheep · 2026-04-20T09:34:17Z

OK, so basically torch doesn’t actually support array API yet (see #159 (comment)), so all the fallback code is just for torch. I don’t think torch should be part of this PR then, we should think about it separately if we want to support it.

So please

move the torch parts (i.e. the fallback code and the torch tests) out if this PR
so wherever there is only a if array_api_compat.is_array_api_obj branch and no @register(... | HasArrayNamespace) branch, create that register
revert the Any changes. I can fix the types but essentially unions will gain a | HasArrayNamespace member and all touched functions gain an overload like def func[A: HasArrayNamespace](x: A, ...) -> A: ...

Regarding the benchmarks: I’m not so sure if this is just static overhead … I think your xp.pow(xp.astype(x, dtype), n) is just genuinely slower than the branch numpy code took before. Can you improve that?

…k510/fast-array-utils into array-api-implementation

for more information, see https://pre-commit.ci

…k510/fast-array-utils into array-api-implementation

for more information, see https://pre-commit.ci

amalia-k510 · 2026-04-27T11:23:04Z

I am still running into mypy issues. It's failing on import jax since jax isn't in the mypy environment. Also, from what I understand, I can't add a mypy override either since the config was removed from main. What would be the best way to handle it? @flying-sheep

flying-sheep

awesome, this is looking very close to perfect! I’ll take a look at the type issues

I can't add a mypy override either since the config was removed from main

what do you mean? the mypy dependencies are directly in .pre-commit-config.yaml (sadly)

flying-sheep · 2026-04-27T11:38:03Z

+) -> NDArray[Any] | np.number[Any]:
+    del keep_cupy_as_array
+    arr = getattr(np, op)(x, axis=axis, **_dtype_kw(dtype, op))
+    return arr.toarray() if isinstance(arr, types.CupyCOOMatrix) else arr  # type: ignore[return-value]


if isinstance(arr, types.CupyCOOMatrix) can’t be true here, right?

also what slowdown do you mean?

numpy has __array_namespace__, so without the explicit np.ndarray registration it'd match HasArrayNamespace and go through array_api_compat.array_namespace() instead of just getattr(np, op)(...). That's the extra array_api_compat overhead showed up in the benchmarks.

As for mypy, sorry, I meant the mypy dependency group was removed, right? So jax isn't installed when mypy runs. Can I just add a [[tool.mypy.overrides]] for jax to the existing [tool.mypy] config?

nvm, just saw the edited comment.

flying-sheep · 2026-04-27T11:39:26Z

    if TYPE_CHECKING:
        assert not isinstance(x, types.DaskArray | types.CSBase | types.CupyCSMatrix)
+    return x**n if dtype is None else np.power(x, n, dtype=dtype)  # type: ignore[operator]


maybe just raise NotImplementedError here?

…k510/fast-array-utils into array-api-implementation

for more information, see https://pre-commit.ci

…k510/fast-array-utils into array-api-implementation

flying-sheep · 2026-04-27T12:03:34Z

+[[tool.mypy.overrides]]
+module = [ "jax", "jax.*" ]
+ignore_missing_imports = true
+


no, we want to typecheck it. I added it to the type check deps

Suggested change

[[tool.mypy.overrides]]

module = [ "jax", "jax.*" ]

ignore_missing_imports = true

flying-sheep · 2026-04-27T12:04:30Z

OK, I started with the typing. As you can see, adding an overload to sum makes sum work where it’s called. This should be done everywhere. I can do it, as you wish.

amalia-k510 · 2026-04-27T13:19:33Z

OK, I started with the typing. As you can see, adding an overload to sum makes sum work where it’s called. This should be done everywhere. I can do it, as you wish.

Thanks for the example! I'd actually like to take that on myself if that's okay.

amalia-k510 and others added 2 commits March 23, 2026 11:28

array-api initially implementation

64f5304

[pre-commit.ci] auto fixes from pre-commit.com hooks

5373050

for more information, see https://pre-commit.ci

flying-sheep reviewed Mar 23, 2026

View reviewed changes

amalia-k510 and others added 7 commits April 15, 2026 16:16

updates in regards to the handler and some array_api handling fixes

f016c39

[pre-commit.ci] auto fixes from pre-commit.com hooks

cca3ad6

for more information, see https://pre-commit.ci

Issues with jax test are fixed, introduced similar tests with pytorch

c1bb155

Merge branch 'array-api-implementation' of https://github.com/amalia-…

c9a8f85

…k510/fast-array-utils into array-api-implementation

pre-commit fixes

bd08c2e

mipy issues fix

b2e3f9b

speed up fix

65e83a6

flying-sheep reviewed Apr 17, 2026

View reviewed changes

amalia-k510 and others added 4 commits April 17, 2026 10:59

Update src/fast_array_utils/stats/_mean_var.py

2a1924d

Co-authored-by: Philipp A. <flying-sheep@web.de>

Addressed the comments

9684213

[pre-commit.ci] auto fixes from pre-commit.com hooks

1e3c296

for more information, see https://pre-commit.ci

chore: simplify

63c5e16

amalia-k510 and others added 12 commits April 27, 2026 12:31

addressing comments about removing is_array_api_obj check

9c8466a

Merge branch 'array-api-implementation' of https://github.com/amalia-…

86a7503

…k510/fast-array-utils into array-api-implementation

import fix

e28f176

ignore comments update and mypy test

c704259

Merge branch 'main' into array-api-implementation

ab6e200

[pre-commit.ci] auto fixes from pre-commit.com hooks

aaccbda

for more information, see https://pre-commit.ci

main version

7e5102a

[pre-commit.ci] auto fixes from pre-commit.com hooks

eed16a9

for more information, see https://pre-commit.ci

residues ignore comments removed

567107a

Merge branch 'array-api-implementation' of https://github.com/amalia-…

1e41d24

…k510/fast-array-utils into array-api-implementation

pyproject, jax optional dependencies

aab50f7

[pre-commit.ci] auto fixes from pre-commit.com hooks

91ef896

for more information, see https://pre-commit.ci

flying-sheep reviewed Apr 27, 2026

View reviewed changes

amalia-k510 and others added 6 commits April 27, 2026 13:55

commented addressed, mypy try again

0f62abc

Merge branch 'array-api-implementation' of https://github.com/amalia-…

d642668

…k510/fast-array-utils into array-api-implementation

[pre-commit.ci] auto fixes from pre-commit.com hooks

37a6634

for more information, see https://pre-commit.ci

mypy comment add

bab4392

Merge branch 'array-api-implementation' of https://github.com/amalia-…

d22f74b

…k510/fast-array-utils into array-api-implementation

types

3d4ee3a

flying-sheep reviewed Apr 27, 2026

View reviewed changes

	[[tool.mypy.overrides]]
	module = [ "jax", "jax.*" ]
	ignore_missing_imports = true

Conversation

amalia-k510 commented Mar 23, 2026

Uh oh!

codecov Bot commented Mar 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

codspeed-hq Bot commented Mar 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Merging this PR will degrade performance by 13.15%

Performance Changes

Uh oh!

flying-sheep left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

amalia-k510 Apr 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

flying-sheep left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

flying-sheep commented Apr 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

amalia-k510 commented Apr 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

flying-sheep left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

amalia-k510 Apr 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

amalia-k510 Apr 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

flying-sheep commented Apr 27, 2026

Uh oh!

amalia-k510 commented Apr 27, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

codecov Bot commented Mar 23, 2026 •

edited

Loading

codspeed-hq Bot commented Mar 23, 2026 •

edited

Loading

amalia-k510 Apr 20, 2026 •

edited

Loading

flying-sheep left a comment •

edited

Loading

flying-sheep commented Apr 20, 2026 •

edited

Loading

amalia-k510 commented Apr 27, 2026 •

edited

Loading

flying-sheep left a comment •

edited

Loading

amalia-k510 Apr 27, 2026 •

edited

Loading

amalia-k510 Apr 27, 2026 •

edited

Loading