Use pytorch as backend for xarrays

I would be interested in using pytorch as a backend for xarrays - because:
a) pytorch is very similar to numpy - so the conceptual overhead is small
b) [most helpful] enable having a GPU as the underlying hardware for compute - which would provide non-trivial speed up
c) it would allow seamless integration with deep-learning algorithms and techniques

Any thoughts on what the interest for such a feature might be ? I would be open to implementing parts of it - so any suggestions on where I could start ?

Thanks

@rgommers

If pytorch implements overrides of NumPy's API via the __array_function__ protocol, then this could work with minimal effort. We are already using this to support sparse arrays (this isn't an official release yet, but functionality is working in the development version).

I think there has been some discussion about this, but I don't know the current status (CC @rgommers). The biggest challenge for pytorch would be defining the translation layer that implements NumPy's API.

Personally, I think the most viable way to achieve seamless integration with deep learning libraries would be to support integration with JAX, which already implements NumPy's API almost exactly. I have an experimental pull request adding __array_function__ to JAX, but it still needs a bit of work to finish it up, e.g., we probably want to hide this behind a flag at first.

@rgommers

I think there has been some discussion about this, but I don't know the current status (CC @rgommers).

The PyTorch team is definitely receptive to the idea of adding __array_function__ and __array_ufunc__, as well as expanding the API for better NumPy compatibility.

Also, they want a Tensor.__torch_function__ styled after __array_function__ so they can make their own API overridable.

The tracking issue for all of this is pytorch/pytorch#22402

The biggest challenge for pytorch would be defining the translation layer that implements NumPy's API.

Agreed. No one is working on __array_function__ at the moment. Implementing it has some backwards compat concerns as well, because people may be relying on np.somefunc(some_torch_tensor) to be coerced to ndarray. It's not a small project, but implementing a prototype with a few function in the torch namespace that are not exactly matching the NumPy API would be a useful way to start pushing this forward.

Personally, I think the most viable way to achieve seamless integration with deep learning libraries would be to support integration with JAX, which already implements NumPy's API almost exactly.

Less familiar with that, but pytorch does have experimental XLA support, so that's a start.

Implementing it has some backwards compat concerns as well, because people may be relying on np.somefunc(some_torch_tensor) to be coerced to ndarray.

Yes, this is a concern for JAX as well. This is a definite downside of reusing NumPy's existing namespace.

It turns out even xarray was relying on this behavior with dask in at least one edge case: #3215

This is a definite downside of reusing NumPy's existing namespace.

We didn't discuss an alternative very explicitly I think, but at least we'll have wide adoption fast. Hopefully the pain is limited ....

I haven't used JAX - but was just browsing through its documentation and it looks super cool. Any ideas on how it compares with Pytorch in terms of:

a) Cxecution speed, esp. on GPU
b) Memory management on GPUs. Pytorch has the 'Dataloader/Dataset' paradigm which uses background multithreading to shuttle batches of data back and forth - along with a lot of tips and tricks on efficient memory usage.
c) support for deep-learning optimization algorithms ?

Within a jit compiled function, JAX's execution speed should be quite competitive on GPUs. It uses the XLA compiler, which was recently enabled by default in TensorFlow.

For data loading and deep learning algorithms, take a look at the examples in the notebooks directory in the JAX repo. The APIs for deep learning in JAX are still undergoing rapid development, so APIs are not quite as stable/usable as pytorch or keras yet, but they are quite capable. See jax.experimental.stax and tensor2tensor.trax for examples.

While it is pretty straightforward to implement a lot of standard xarray operations with a pytorch / Jax backend (since they just fallback on native functions) - it will be interesting to think about how to implement rolling operations / expanding / exponential window in a way that is both efficient and maintains differentiability.

Expanding and exponential window operations would be easy to do leveraging RNN semantics - but doing rolling using convolutions is going to be very inefficient.

Do you have any thoughts on this?

I have not thought too much about these yet. But I agree that they will probably require backend specific logic to do efficiently.

…

On Fri, Aug 23, 2019 at 12:13 PM firdaus janoos ***@***.***> wrote: While it is pretty straightforward to implement a lot of standard xarray operations with a pytorch / Jax backend (since they just fallback on native functions) - it will be interesting to think about how to implement rolling operations / expanding / exponential window in a way that is both efficient and maintains differentiability. Expanding and exponential window operations would be easy to do leveraging RNN semantics - but doing rolling using convolutions is going to be very inefficient. Do you have any thoughts on this? — You are receiving this because you commented. Reply to this email directly, view it on GitHub <#3232?email_source=notifications&email_token=AAJJFVWRVLTFNT3DYOZIJB3QGASFBA5CNFSM4ING6FH2YY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOD5A6IWY#issuecomment-524411995>, or mute the thread <https://github.com/notifications/unsubscribe-auth/AAJJFVQ7JBUNO3CAIFGVJ63QGASFBANCNFSM4ING6FHQ> .

@mrocklin

This might be a good time to revive this thread and see if there is wider interest (and bandwidth) in having xarray use CuPy (https://cupy.chainer.org/ ) as a backend (along with numpy). It appears to be a plug-and-play replacement for numpy - so it might not have all the issues that were brought up regarding pytorch/jax ?

Any thoughts ?
cc @mrocklin

Just chiming in quickly. I think there's definitely interest in doing this through NEP-18.

It looks like CUDA has implemented __array_function__ (https://docs-cupy.chainer.org/en/stable/reference/interoperability.html) so many things may "just work". There was some work earlier on plugging in pydata/sparse, and there is some ongoing work to plug in pint. With both these efforts, a lot of xarray's code should be "backend-agnostic" but its not perfect.

Have you tried creating DataArrays with cupy arrays yet? I would just try things and see what works vs what doesn't.

Practically, our approach so far has been to add a number of xfailed tests (test_sparse.py and test_units.py) and slowly start fixing them. So that's one way to proceed if you're up for it.

@jacobtomlinson

@jacobtomlinson gave CuPy a go a few months back. I seem to remember that he ran into a few problems but it would be good to get those documented here.

Yeah Jacob and I played with this a few months back. There were some issues, but my recollection is pretty hazy. If someone gives this another try, it would be interesting to hear how things go.

If you have any pointers on how to go about this - I can give it a try.

…

Well here's a blogpost on using Dask + CuPy. Maybe start there and build up to using Xarray.

I started having a look at making xarray work with the array API here: tomwhite@c72a1c4. Some basic operations work (preserving the underlying array): tomwhite@929812a. If there's interest, I'd be happy to turn this into a PR with some tests.

I'd be happy to turn this into a PR with some tests.

Absolutely!

Opened #6804

Glad to see progress on this!! 👏

Just curious though, seeing this comment in the PR:

Note: I haven't actually tested this with pytorch (which is the motivating example for #3232).

Are we sure this closes the issue? And, how can we try it out? Even lacking docs, a comment explaining how to set it up would be great, and I can do some testing on my end. I understand that it's an experimental feature.

@hsharrison

Hi @hsharrison - thanks for offering to do some testing. Here's a little demo script that you could try, by switching numpy.array_api to pytorch: tomwhite@929812a

Nice that it's so simple. I think it can't be tested with pytorch until they compete pytorch/pytorch#58743, right?

Or we should just try passing torch.tensor into xarray directly?

I think it can't be tested with pytorch until they compete pytorch/pytorch#58743, right?

It needs __array_namespace__ to be defined to activate the new code path.

Makes sense, then I'll wait for pytorch/pytorch#58743 to try it.

While it is true to use PyTorch Tensors directly, one would need the Array API implemented in PyTorch. One could use them indirectly by converting them zero-copy to CuPy arrays, which do have Array API support

Even after define __array_namespace__ to torch.Tensor, there could be occasionally RuntimeError trigger for checking hasattr(torch.Tensor, "imag"). Minimal example to reproduce the problem is follows.

# Python 3.11.7, xarray 2025.12.0, Torch 2.9.1, numpy 2.2.6
import array_api_compat.torch as xp_torch
import torch
import xarray as xr

torch.Tensor.__array_namespace__ = lambda self, **kwargs: xp_torch
da = xr.DataArray(torch.zeros, coords=dict(x=[0], y=[1, 2]))
print(f"Underlying type: {type(da.data)}")

print(da.data.dtype)

da.data[0, 0] += 1
print(da)
print(type(da.sum().data))
print(type(da.sel(x=0).data))
print(da.transpose("y", "x"))

It occassionally triggers bug

Traceback (most recent call last):
File "", line 99, in <module>
    print(type(da.sum().data))
File "./.venv/lib/python3.11/site-packages/xarray/core/_aggregations.py", line 3167, in sum
    return self.reduce(
           ^^^^^^^^^^^^
...
File "./.venv/lib/python3.11/site-packages/xarray/namedarray/core.py", line 205, in from_array
    if isinstance(data, _arrayfunction_or_api) and not isinstance(data, np.generic):
File "./lib/python3.11/typing.py", line 2006, in __instancecheck__
    if all(hasattr(instance, attr) and
RuntimeError: imag is not implemented for tensors with non-complex dtypes.

The randomness is from _get_protocol_attrs(cls) in typing.py is using a set. Thus torch.Tensor.imag can raise exception if imag is checked first. A quick work around is to overwrite torch.Tensor.imag to allow it return None for non-complex tensor. Another way is to switch check order in xarray at xarray/namedarray/_typing.py to check array api first

_arrayfunction_or_api = (_arrayfunction, _arrayapi)

Can you open another issue for this one? Pytorch's behavior is arguably a bit pathological here, but the array API compatible strategy for Xarray of using xp.imag() instead of the .imag property should be straightforward.

…

On Wed, Dec 31, 2025 at 12:37 PM cshuphys ***@***.***> wrote: *cshuphys* left a comment (pydata/xarray#3232) <#3232 (comment)> Even after define __array_namespace__ to torch.Tensor, there could be occasionally RuntimeError trigger for checking hasattr(torch.Tensor, "imag"). Minimal example to reproduce the problem is follows. # Python 3.11.7, xarray 2025.12.0, Torch 2.9.1, numpy 2.2.6 import array_api_compat.torch as xp_torch import torch torch.Tensor.__array_namespace__ = lambda self, **kwargs: xp_torch da = tensor.Tensor(torch.zeros, coords=dict(x=[0], y=[1, 2])) print(f"Underlying type: {type(da.data)}") print(da.data.dtype) da.data[0, 0] += 1 print(da) print(type(da.sum().data)) print(type(da.sel(x=0).data)) print(da.transpose("y", "x")) It occassionally triggers bug Traceback (most recent call last): File "", line 99, in <module> print(type(da.sum().data)) File "./.venv/lib/python3.11/site-packages/xarray/core/_aggregations.py", line 3167, in sum return self.reduce( ^^^^^^^^^^^^ ... File "./.venv/lib/python3.11/site-packages/xarray/namedarray/core.py", line 205, in from_array if isinstance(data, _arrayfunction_or_api) and not isinstance(data, np.generic): File "./lib/python3.11/typing.py", line 2006, in __instancecheck__ if all(hasattr(instance, attr) and RuntimeError: imag is not implemented for tensors with non-complex dtypes. The randomness is from _get_protocol_attrs(cls) in typing.py is using a set. Thus torch.Tensor.imag can raise exception if imag is checked first. A quick work around is to overwrite torch.Tensor.imag to allow it return None for non-complex tensor. Another way is to switch check order in xarray at xarray/namedarray/_typing.py to check array api first _arrayfunction_or_api = (_arrayfunction, _arrayapi) — Reply to this email directly, view it on GitHub <#3232 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJJFVUV44RQAWCOVKI2S4L4EQXZTAVCNFSM6AAAAACQNJEEYKVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZTOMBSHA3DSNRXGQ> . You are receiving this because you modified the open/close state.Message ID: ***@***.***>

Navigation Menu

Repository navigation

Use pytorch as backend for xarrays #3232

25 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Use pytorch as backend for xarrays #3232

Description

Activity

shoyer commented on Aug 20, 2019

rgommers commented on Aug 20, 2019

rgommers commented on Aug 20, 2019

shoyer commented on Aug 20, 2019

rgommers commented on Aug 20, 2019

fjanoos commented on Aug 23, 2019

shoyer commented on Aug 23, 2019

fjanoos commented on Aug 23, 2019

shoyer commented on Aug 23, 2019

fjanoos commented on Mar 30, 2020

dcherian commented on Mar 30, 2020

jhamman commented on Mar 30, 2020

jakirkham commented on Mar 30, 2020

fjanoos commented on Mar 31, 2020

jakirkham commented on Mar 31, 2020

25 remaining items

tomwhite commented on Jul 13, 2022

dcherian commented on Jul 13, 2022

tomwhite commented on Jul 18, 2022

hsharrison commented on Jul 20, 2022

tomwhite commented on Jul 20, 2022

hsharrison commented on Jul 20, 2022

tomwhite commented on Jul 20, 2022

hsharrison commented on Jul 20, 2022

jakirkham commented on Jul 20, 2022

cshuphys commented on Dec 31, 2025

shoyer commented on Jan 1, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions