Combining JAX and xarray for large scale computing

[Aceticia]

I'm working on computational neuroscience and I do a lot of large scale computations (perhaps not as big as genomics folks). These computations are usually simple linear algebra manipulations, parallelized to a bunch of CPUs. Meanwhile, our lab have a large number of GPUs for use. What would be a good way to take advantage of the easy usage of xarray and JAX?

In my workflow, I already have a lot of vectorized calls and it's saving me a lot of time. But these are still a large number of small linear algebra operations that might see huge gains when performed on GPUs.

A naive way is to simply convert to jax in a vectorized function. For example, I vectorize the calls to a number of brain scans. [num_subjects, num_conditions, x, y, z, features] where xyz is the spatial dimensions. I have a function func that processes certain features. I apply ufunc to vectorize over the first 5 axes. Then I just make func jax, and don't touch anything else.

But this is likely going to incur a large conversion overhead. In the case that entire data fits into memory un-chunked, what do you think is a better approach?

[keewis]

if you repeatedly would be converting to / from jax, it would probably be best to keep everything in jax and have xarray wrap that.

According to jax-ml/jax#17107, support for doing so seems to be close but not quite there yet. In particular we'd need jax to implement the array API (i.e., define __array_namespace__) – see jax-ml/jax#16099 a in-progress PR.

[ConstantinVasilev]

Hi @keewis, the Jax PR (jax-ml/jax#16099) is merged, so is there an updated answer to this discussion?

[keewis]

nothing that I'm aware of, but I'm not using JAX myself so I'm not the best person to ask.

However, since JAX appears to have Array API compatibility now, I'd invite you to just try it out: create a JAX array and wrap it using a Dataset or DataArray object. If that works, try calling a couple of methods and see what fails, and in any case I'd encourage you to report back here on how it went (also, feel free to ask any questions regarding xarray usage here).

[zazbone]

Hi !
I'm in a similar situation where I use Jax as computation backend, but I could also benefit from xarray Dataset
I've tried to init a DataArray with as Jax array as input data, but it looks like xarray does an automatic conversion to ndarray
I raise the subject again to know if there is update or plan on jax or array api standard support

Cheers

[keewis]

looks like you have to use jax.experimental.array_api instead of jax.numpy for this to work:

In [5]: import jax.experimental.array_api as jxp
   ...: import xarray as xr
   ...: 
   ...: arr = xr.DataArray(jxp.zeros((5, 3)), dims=["y", "x"])
   ...: arr.sum(dim="y")
Out[5]: 
<xarray.DataArray (x: 3)> Size: 12B
Array([0., 0., 0.], dtype=float32)
Dimensions without coordinates: x

(note the Array instead of array)

which doesn't mean all xarray operations are supported (see the list of known issues, and the rest you'll have to check yourself), but at least standard reductions appear to have worked fine.

[TomNicholas]

Wait this is big! I didn't know jax.experimental.array_api existed!

Now would be a really good time to finish off that framework for automatically testing different duck-typed arrays in xarray... 😅

[keewis]

that will take a lot of concentrated effort... we can try doing some kind of on-call sprint, if you'd like (if so, to be discussed somewhere else).

[TomNicholas]

I would be down for that actually. I also met a bunch of people at SciPy who are interested in ragged arrays in xarray. But I am supposed to be on holiday right now, so maybe remind me sometime after next week? 😁

[zazbone]

Ho, good to see it work, I will wait until Jax stabilized it but it's cool
Thanks

[allen-adastra]

Starting from JAX 0.4.32, we should be able to wrap jax.Array with xarray structures.

Two additional issues for Xarray + JAX integration have been identified:
1. Automatic dispatching based on array type #7848
2. Registering Xarray structures as PyTree nodes

Issue 2 has an example solution in the GraphCast project.

The main challenge I've run into there is that many JAX libraries (e.g. Diffrax) use tree manipulation patterns that produce "dummy trees" like boolean masks. If a xr.DataArray has its data attribute replaced with a True, it will trigger all kinds of errors.

I can think of three solutions to this:

1. Avoid inputting xarray types into functions that do this kind of masking
2. Introduce flags into xarray to disable checks to allow for things like tree masks
3. Implement flatten functions that convert xarray types into another type (e.g. XJDataArray) and implement to_xarray() methods to allow the user to get back the original structure inside jitted functions (experimental implementation of this approach)

[allen-adastra]

Something else I've run into: xr.IndexVariable converts JAX arrays to numpy arrays. Is this intended? Seems like it makes sense but just want to check.

@given(var=index_variables, ufunc=ufunc_strat)
@settings(deadline=None)
def test_index_variables(var: xr.Variable, ufunc):
    var = var.to_index_variable()

    # Index variables should not be wrapped in a jax array.
    assert not isinstance(var.data, jax.Array)

[TomNicholas]

I've got a test with some hypothesis strategies for types,

Yay someone's actually using that!

Do we have a xr.DataArray or xr.Dataset strategy btw?

No, but there is a 90% finished PR to add them... (#6908)

[allen-adastra]

Ah cool! Would be cool to have that (not that you don't already have a lot going on here :>))

[shoyer]

Something else I've run into: xr.IndexVariable converts JAX arrays to numpy arrays. Is this intended? Seems like it makes sense but just want to check.

Yes, this is intended. Variables with an index get loaded into memory, with an associated pandas.Index. They should be treated as static information as far as JAX is concerned.

[allen-adastra]

@allen-adastra In your failing example with Dataset, you are iterating over Dataset.data_vars. I'm not sure we can reasonably expect that or any other transformation that creating invalid xarray data values to work. What should work are cases where JAX uses pytrees for internal book-keeping but doesn't surface those in external APIs.

I would love to see a specific list of JAX transformations (including external ones like Equinox and NNX) that do or do not work.

Merged a PR that does checks some basic JAX and Equinox transforms, and things seem to work pretty well as long as the shape of the underlying array doesn't change (including jax.grad, although you gotta make sure you expose the underlying array with .data on the function output).

I saw @mjwillson had a particular solution for handling dim changes in the GraphCast project, but that requires the user to manually specify a context. Would be nice if there was some way to add dummy dimensions automatically.

[allen-adastra]

I pushed a v0.0.1 solution to PyPi with an acknowledgement of the folks in this thread. Would welcome testers to find new and exciting ways to break it :)