ENH: speed up array_namespace

[crusaderky]

Related:

• ENH: cache helper functions #308
• ENH: Speed up array_namespace through caching scipy/scipy#23057

Tested locally vs. jax==0.4.31.

Benchmark

>>> import importlib >>> from array_api_compat import array_namespace >>> for xp_name in ("numpy", "cupy", "dask.array", "torch", "jax.numpy", "sparse", "ndonnx"): ... print(xp_name) ... xp = importlib.import_module(xp_name) ... a = xp.asarray(1) ... %timeit array_namespace(a)

Before

numpy 1.31 μs ± 19.9 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each) cupy 1.36 μs ± 5.16 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each) dask.array 1.77 μs ± 33.6 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each) torch 1.52 μs ± 4.8 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each) jax.numpy 1.4 μs ± 19.8 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each) sparse 1.75 μs ± 9.39 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each) ndonnx 1.99 μs ± 10.9 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each) 

After

numpy 673 ns ± 11.8 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each) cupy 494 ns ± 8.45 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each) dask.array 498 ns ± 3.87 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each) torch 510 ns ± 5.91 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each) jax.numpy 521 ns ± 3.06 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each) sparse 766 ns ± 2.13 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each) ndonnx 727 ns ± 9.62 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each) 

Note that the extra slowness of numpy is due to jax-ml/jax#20620

[crusaderky]

Note: this test requires jax>=0.6.1 to pass with api_version=2024.12. I didn't see much of a point adding special handling for older versions.

[ev-br]

Could you please resolve the conflict

[crusaderky]

@ev-br fixed

[ev-br]

Grumble grumble:

========================================================================================= short test summary info ========================================================================================= FAILED tests/test_array_namespace.py::test_array_namespace[jax.numpy-2024.12-False] - ValueError: api_version='2024.12' is not available; available versions are: ['2023.12'] FAILED tests/test_array_namespace.py::test_array_namespace[jax.numpy-2024.12-None] - ValueError: api_version='2024.12' is not available; available versions are: ['2023.12'] ============================================================== 2 failed, 410 passed, 105 skipped, 6 xfailed, 2 xpassed, 46 warnings in 7.66s ============================================================== 

EDIT: sorry, no; my local Jax copy was too old. Nevermind.

[ev-br]

Thanks @crusaderky , merged.