Skip to content

Adding array-api-compat fallback#159

Draft
amalia-k510 wants to merge 2 commits intoscverse:mainfrom
amalia-k510:array-api-implementation
Draft

Adding array-api-compat fallback#159
amalia-k510 wants to merge 2 commits intoscverse:mainfrom
amalia-k510:array-api-implementation

Conversation

@amalia-k510
Copy link

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.

@codecov
Copy link

codecov bot commented Mar 23, 2026

Codecov Report

❌ Patch coverage is 41.66667% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 93.71%. Comparing base (395b37e) to head (5373050).

Files with missing lines Patch % Lines
src/fast_array_utils/stats/_is_constant.py 0.00% 11 Missing ⚠️
src/fast_array_utils/conv/_to_dense.py 40.00% 3 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #159      +/-   ##
==========================================
- Coverage   96.78%   93.71%   -3.07%     
==========================================
  Files          19       19              
  Lines         466      493      +27     
==========================================
+ Hits          451      462      +11     
- Misses         15       31      +16     

☔ 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
Copy link

codspeed-hq bot commented Mar 23, 2026

Merging this PR will degrade performance by 26.04%

❌ 17 regressed benchmarks
✅ 215 untouched benchmarks

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

Performance Changes

Benchmark BASE HEAD Efficiency
test_stats_benchmark[numpy.ndarray-1d-all-int32-mean_var] 120.8 ms 162.1 ms -25.5%
test_stats_benchmark[numpy.ndarray-2d-all-float32-mean_var] 123.4 ms 146.9 ms -15.99%
test_stats_benchmark[numpy.ndarray-2d-all-int32-mean_var] 120.8 ms 162.1 ms -25.5%
test_stats_benchmark[numpy.ndarray-1d-all-float32-mean_var] 123.4 ms 146.9 ms -15.98%
test_stats_benchmark[numpy.ndarray-2d-ax0-float64-mean_var] 170.1 ms 189.2 ms -10.12%
test_stats_benchmark[numpy.ndarray-2d-ax0-float32-mean_var] 120.1 ms 143.6 ms -16.37%
test_stats_benchmark[numpy.ndarray-2d-ax1-float32-mean_var] 123.6 ms 147 ms -15.96%
test_stats_benchmark[numpy.ndarray-2d-ax0-int32-mean_var] 117.5 ms 158.8 ms -26.04%
test_stats_benchmark[numpy.ndarray-2d-ax1-int32-mean_var] 121 ms 162.3 ms -25.47%
test_stats_benchmark[scipy.sparse.csc_array-1d-all-float32-mean_var] 123.5 ms 147 ms -15.96%
test_stats_benchmark[scipy.sparse.csc_array-1d-all-int32-mean_var] 120.9 ms 162.2 ms -25.47%
test_stats_benchmark[scipy.sparse.csc_array-2d-all-float32-mean_var] 123.5 ms 147 ms -15.95%
test_stats_benchmark[scipy.sparse.csc_array-2d-all-int32-mean_var] 120.9 ms 162.2 ms -25.48%
test_stats_benchmark[scipy.sparse.csr_array-1d-all-float32-mean_var] 123.5 ms 147 ms -15.95%
test_stats_benchmark[scipy.sparse.csr_array-2d-all-int32-mean_var] 120.9 ms 162.2 ms -25.47%
test_stats_benchmark[scipy.sparse.csr_array-1d-all-int32-mean_var] 120.9 ms 162.2 ms -25.47%
test_stats_benchmark[scipy.sparse.csr_array-2d-all-float32-mean_var] 123.5 ms 147 ms -15.95%

Comparing amalia-k510:array-api-implementation (5373050) with main (395b37e)

Open in CodSpeed

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant