Add basic JAX support #84

asmeurer · 2024-02-06T22:27:20Z

Unlike other modules, JAX array API support is fully in JAX itself in the jax.experimental.array_api submodule, so the only thing that is done here is to add JAX support to the helper functions. This also means that we do not run array-api-tests on JAX.

This also makes the various is_numpy_array, is_cupy_array, etc. functions public, as I noticed someone was using them on GitHub and they seem like they could be useful.

… public

Unlike other libraries, there is no wrapping for JAX. Actual JAX array_api support is in JAX itself in the jax.experimental.array_api submodule. This just adds JAX support to the various helper functions. This also means that we do not run array-api-tests on JAX. Closes data-apis#83.

asmeurer · 2024-02-06T22:28:43Z

array_api_compat/common/_helpers.py

@jakevdp, I could mostly use your review for the changes in this file.

asmeurer · 2024-02-06T22:29:02Z

array_api_compat/common/_helpers.py

+            # jax.numpy is already an array namespace, but requires this
+            # side-effecting import for __array_namespace__ and some other
+            # things to be defined.
+            import jax.experimental.array_api as jnp


Will this import go away at some point? Should we guard against that?

It's not true that jax.numpy is already an array namespace. jax.experimental.array_api is an array namespace, and we hope to make jax.numpy an array namespace in the future.

OK, I fixed the comment. The question still remains though? Should I add a guard here like

try: import jax.experimental.array_api as jnp except ImportError: import jax.numpy as jnp

for a future JAX version when jax.experimental.array_api gets removed? Or will it never be removed?

Yes, I think that's probably a reasonable way to future-proof this.

The downside of that particular line of code is it will also pass through jax.numpy for older JAX versions that don't have jax.experimental.array_api. I like to avoid explicit version checks in this library if I can, but maybe that's the best thing to do here.

Or maybe we can just change the logic to this once JAX starts to remove (deprecates?) the experimental import.

asmeurer · 2024-02-06T22:29:12Z

array_api_compat/common/_helpers.py

+        # JAX has .device() as a method, but it is being deprecated so that it
+        # can become a property, in accordance with the standard. In order for
+        # this function to not break when JAX makes the flip, we check for
+        # both here.


Does this logic seem OK?

This looks OK, but it will not work with the to_device function as currently defined in jax.experimental.array_api. That expects to be passed the bound method object – it's a hack, and one of the reasons that this is still considered experimental.

The other problem here is that, in general, JAX arrays can live on multiple devices (in which case arr.device() will error). It's not clear to me from the Array API design how multi-device arrays are meant to be supported.

I don't know the answer to that. I would bring it up on the array API repo. https://github.com/data-apis/array-api/. As far as I know it hasn't really been discussed.

It looks like both forms work

>>> import jax.experimental.array_api as xp >>> x = xp.asarray([1, 2, 3]) >>> x.to_device(x.device) Array([1, 2, 3], dtype=int32) >>> x.to_device(x.device()) Array([1, 2, 3], dtype=int32)

asmeurer · 2024-02-06T22:29:26Z

array_api_compat/common/_helpers.py

+    import jax
+    if stream is not None:
+        raise NotImplementedError
+    return jax.device_put(x, device)


Is this helper function for to_device correct?

Yes, device_put would be the right method here.

jakevdp · 2024-02-06T22:34:34Z

array_api_compat/common/_helpers.py

+
+    import jax.numpy
+
+    return isinstance(x, jax.numpy.ndarray)


isinstance(x, jax.Array) would be more concise.

asmeurer · 2024-02-06T22:34:40Z

Oh we should also probably add a basic test for the device helpers at https://github.com/data-apis/array-api-compat/blob/main/tests/test_common.py

asmeurer · 2024-02-06T23:32:35Z

Apparently we've been supporting "cpu" as a special host device here (#40). This is still being discussed for the standard data-apis/array-api#626.

I'm actually unsure if it's a good idea for us to be supporting that here given that we haven't really agreed about it in the standard. But if we did want to support it for JAX, how would we? I don't see how to actually access jax.CpuDevice.

jakevdp · 2024-02-06T23:34:49Z

jax.devices('cpu')[0] would return the first CPU device (if available).

jakevdp

I think this looks good!

jakevdp · 2024-02-07T00:03:55Z

array_api_compat/common/_helpers.py


-    return isinstance(x, jax.numpy.ndarray)
+    return isinstance(x, jax.Array)


Should this be more guarded? e.g. what if someone has a module named jax.py in their path, it could lead to an AttributeError here and make array_api_compat unusable.

I don't know. That sort of thing tends to just break everything anyway. I've never really felt that libraries should protect against that sort of thing.

Anyway, the whole point of this function is to be guarded. It won't import jax unless it's already been imported, because jax is a slow import and unnecessary if a user is using another array library.

More precisely, it won't import jax unless something named jax has been imported, and then array_api_compat will irrecoverably error if that jax module isn't the one we're expecting, or is an older version that doesn't define the attributes we reference.

A similar issue exists for every other package name referenced in this module.

My feeling is: it costs virtually nothing to wrap this all in an appropriate try ... except, so why not do so?

I'm still not so sure this is a good idea. Usually if you have that sort of thing it will be an error for a lot of things, not just array_api_compat. My worry here is that guarding isn't as straightforward as it might seem. Wrapping everything in try/except could mean we end up silencing legitimate errors.

OK, sounds good.

asmeurer · 2024-02-08T20:09:44Z

Removed the "cpu" device logic. It looks like that isn't going to be part of the standard. I opened #86 about removing it for cupy as well.

test_helpers.py is too confusing alongside tests/_helpers.py.

This requires using subprocess to test that it works even if the side-effecting jax.experimental.array_api hasn't been imported yet.

asmeurer added 4 commits February 6, 2024 14:58

Make is_numpy_array, is_cupy_array, is_torch_array, and is_dask_array…

2ee6902

… public

Note dask in a couple of places in the README

12b5294

Install JAX on CI

583f6bb

asmeurer mentioned this pull request Feb 6, 2024

JAX support #83

Closed

asmeurer commented Feb 6, 2024

View reviewed changes

array_api_compat/common/_helpers.py

Copy link

Member Author

asmeurer Feb 6, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@jakevdp, I could mostly use your review for the changes in this file.

asmeurer commented Feb 6, 2024

View reviewed changes

asmeurer mentioned this pull request Feb 6, 2024

Add dask to array-api-compat #76

Merged

Install JAX as jax[cpu]

9c8bed6

jakevdp reviewed Feb 6, 2024

View reviewed changes

asmeurer added 2 commits February 6, 2024 15:36

Fix a comment

6d59ae8

Use jax.Array instead of jax.numpy.ndarray in is_jax_array

ce07cd9

asmeurer added 5 commits February 6, 2024 16:40

Use the native jax to_device() method

ddb313e

Add a basic test for device() and to_device()

6004b97

XFAIL the device test for dask.array

701a5ef

Allow to_device(x, "cpu") for JAX arrays

049d557

Move tests in test_common.py to test_helpers.py

db667ea

asmeurer mentioned this pull request Feb 6, 2024

Add ruff to ci setup #82

Merged

jakevdp reviewed Feb 7, 2024

View reviewed changes

Remove "cpu" device from jax to_device()

aafbbaa

asmeurer added 4 commits February 8, 2024 13:21

Merge branch 'main' into jax

919ec41

Rename import_or_skip_cupy to import_ and move jax logic into it

fa758f7

Skip JAX tests in Python 3.8

6c338ca

Rename test_helpers.py to test_common.py

244462f

test_helpers.py is too confusing alongside tests/_helpers.py.

asmeurer added 3 commits February 8, 2024 13:29

Fix ruff warnings

bff9bf2

Add jax.numpy to the test_array_namespace()

264e6c3

This requires using subprocess to test that it works even if the side-effecting jax.experimental.array_api hasn't been imported yet.

Don't run the jax.numpy array_namespace test in Python 3.8

e7aff0f

asmeurer merged commit 645f9a8 into data-apis:main Feb 9, 2024
9 of 27 checks passed

adityagoel4512 mentioned this pull request Jun 17, 2024

Add ndonnx #150

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add basic JAX support #84

Add basic JAX support #84

asmeurer commented Feb 6, 2024

asmeurer Feb 6, 2024

asmeurer Feb 6, 2024

jakevdp Feb 6, 2024

asmeurer Feb 6, 2024

jakevdp Feb 6, 2024

asmeurer Feb 7, 2024

asmeurer Feb 6, 2024

jakevdp Feb 6, 2024

jakevdp Feb 6, 2024 •

edited

Loading

asmeurer Feb 6, 2024

asmeurer Feb 6, 2024

asmeurer Feb 6, 2024

jakevdp Feb 6, 2024

jakevdp Feb 6, 2024

asmeurer commented Feb 6, 2024

asmeurer commented Feb 6, 2024

jakevdp commented Feb 6, 2024 •

edited

Loading

jakevdp left a comment

jakevdp Feb 7, 2024

asmeurer Feb 7, 2024

jakevdp Feb 7, 2024 •

edited

Loading

asmeurer Feb 8, 2024

jakevdp Feb 8, 2024

asmeurer commented Feb 8, 2024


		return isinstance(x, jax.numpy.ndarray)
		return isinstance(x, jax.Array)

Add basic JAX support #84

Add basic JAX support #84

Conversation

asmeurer commented Feb 6, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jakevdp Feb 6, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

asmeurer commented Feb 6, 2024

asmeurer commented Feb 6, 2024

jakevdp commented Feb 6, 2024 • edited Loading

jakevdp left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jakevdp Feb 7, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

asmeurer commented Feb 8, 2024

jakevdp Feb 6, 2024 •

edited

Loading

jakevdp commented Feb 6, 2024 •

edited

Loading

jakevdp Feb 7, 2024 •

edited

Loading