[smaddox]

I'm curious, have you considered using PyTorch or JAX for tensor processing? ML libraries seem to be much further along when it comes to performing compute-intensive, hardware-accelerated operations on Tensor's. And you get gradients basically for free (in terms of developer time). Also, the kernel compiler being added PyTorch 2 looks very promising.

[counters]

The primary issue in this domain is not compute - it's I/O, especially when you need to perform complex computations with intermediate data that doesn't fit into memory.

[rabernat]

PyTorch and JAX are used heavily in climate science on the ML side. For more general analytics, not so much. Many of our users like to use Xarray as a high-level API. There has been some work to integrate Xarray with PyTorch (https://github.com/pydata/xarray/issues/3232) but we're not there yet.

The Python Array API standard should help align these different back-ends: https://data-apis.org/array-api/latest/