ENH: implement at

[crusaderky]

Implement a new at(x, idx) or at(x)[idx] function, mocking the syntax of JAX's omonymous method .
This is propaedeutic to JAX support in libraries that support the Array API, e.g. scipy.

Moved from data-apis/array-api-compat#205

Blockers

• ENH: Test for read-only arrays array-api-compat#205
• BUG: cp.generic is an alias for np.generic array-api-compat#207
• BUG: dask.array: asarray(..., copy=None) copies in dask==2024.12 array-api-compat#211

[crusaderky]

Rebased on top of #58

[crusaderky]

• Reverted rebase onto Saner linters #58
• Fixed all linters issues
• Added unit tests for all backends
• Removed **kwargs. Their main original appeal was allowing to pass performance tweaks to JAX, which are indices_are_sorted=True and unique_indices=True. However, I have now realised that the array API specification does not support integer array indices, at all, and both of these parameters are pointless without.
• I'm having big second thoughts on get() - read below.

[crusaderky]

Any idea how to fix this? Looks like the ubuntu_latest VM has an obsolete driver (or more likely no driver)

ERROR cupy_backends.cuda.api.runtime.CUDARuntimeError: cudaErrorInsufficientDriver: CUDA driver version is insufficient for CUDA runtime version

[lucascolley]

I don't think we can get GPU CI without paying someone for it, cc @rgommers .

[crusaderky]

I don't think we can get GPU CI without paying someone for it, cc @rgommers .

Can cupy run on a CPU-only host?
I've disabled it in CI for the time being

[rgommers]

I don't think we can get GPU CI without paying someone for it, cc @rgommers .

Yep, that's on my radar to push forward this month, on multiple projects. Please feel free to open a new issue and assign it to me. I think we can hook up a shared GPU runner between this project and array-api-compat.

Can cupy run on a CPU-only host?

I don't think so.

[lucascolley]

I resolved the merge conflicts after adding numpydoc to pre-commit and made some minor tweaks to the docs.

[crusaderky]

@lucascolley if you don't mind building against array-api-compat git tip until their next release, this PR is ready for final review and merge

[lucascolley]

@lucascolley if you don't mind building against array-api-compat git tip until their next release, this PR is ready for final review and merge

That's fine with me, but might be cleaner to get an array-api-compat release out first unless there are any big blockers.

Let's give @rgommers and @jakevdp time to take a look if they would like to.

Once (or before) this is merged, do you think you could make a PR to my branch at scikit-learn/scikit-learn#30340? I think we will have to transition scikit-learn from using array-api-compat as an optional dependency to vendoring it, but that sounds feasible based on scikit-learn/scikit-learn#30367 (comment).

[crusaderky]

Once (or before) this is merged, do you think you could make a PR to my branch at scikit-learn/scikit-learn#30340? I think we will have to transition scikit-learn from using array-api-compat as an optional dependency to vendoring it, but that sounds feasible based on scikit-learn/scikit-learn#30367 (comment).

lucascolley/scikit-learn#2

[crusaderky]

@lucascolley ready for merge and release! 🥳

[lucascolley]

good to proceed @rgommers ?