Array API backend

[tomwhite]

Fixes #315

This uses https://github.com/data-apis/array-api-compat which makes NumPy conform to the array API. There are a couple of places where we use functions not in the array API (which fall back to NumPy): take_along_axis for arg reductions, and nan functions.

[alxmrs]

Why not simply import namespace as np? This would minimize the diff.

[tomwhite]

It would, but I wanted to signal that this may not be regular NumPy. By using nxp here, it makes it easy to search for np to find places that are still using regular NumPy.

[alxmrs]

I'm excited to see this as simple place to drop in a Jax array implementation :) Great work.

[tomwhite]

It will be interesting to see if this works with the recent work in JAX on the Array API.

[tomwhite]

Thanks for taking a look at this @alxmrs!

This has conflicts that I need to resolve, but it also conflicts with the work I've been doing for TensorStore in #187, which is why I opened #322. But maybe it's OK to merge this, and figure out the TensorStore part later?

[TomNicholas]

This uses https://github.com/data-apis/array-api-compat which makes NumPy conform to the array API.

You might be able to avoid the extra dependency on array-api-compat now that numpy v1.26.0 introduced numpy.array_api.

[TomNicholas]

But maybe it's OK to merge this, and figure out the TensorStore part later?

I vote merge this and think first about generalizing the array type in cases that don't also require changing the storage layer (e.g. cupy/sparse/pint).

[tomwhite]

This uses https://github.com/data-apis/array-api-compat which makes NumPy conform to the array API.

You might be able to avoid the extra dependency on array-api-compat now that numpy v1.26.0 introduced numpy.array_api.

NumPy 1.22 introduced numpy.array_api, 1.26 updates it to the latest spec. The reason for using the compat layer is that it falls back to some functions that we use that are not in the standard yet (like take_along_axis).

If we're worried about the extra dependency then array-api-compat can be vendored, but I'm not sure that's necessary.