Array API: remove array mutation from functions/write JAX version of each array mutating function

[Saransh-cpp]

Add your issue here

Conversation from Slack -

2 quick things from my initial experiments -

We will need custom JAX implementations for most of the functions as JAX arrays are immutable. See how -

• s2fft (each file has a corresponding *_jax file - Matt Price worked on this), and
• awkward (all other backends have common reducers but the JAX backend has specialized reducers - I worked on this)
  were developed independently, but they both have one common pattern, a standard implementation of a particular function and then a JAX implementation of the same function.

Given that the project's first half is more focussed on GPUs and less on AD, I experimented with CuPy, but their Array API support is still experimental (and I might have also stumbled on a bug - cupy/cupy#8747).

I thought the only place in the code where we cannot trivially get around the mutation was in iternorm() but I haven't actually gone through and checked

Oh yes, or remove all the lines of code that are mutating arrays. Maybe I should check for all such lines using ruff.

So the mutation in iternorm would be quite organically fixed by the "deep dive" project outside the GPU work that I was mentioning. We definitely have a path forward there.

My inclination would be that we change all functions to not mutate arrays and then we can use jax and everything else in one go. Is there any disadvantage of this approach?

Only in the way iternorm currently works, since it always needs to discard a large array and reallocate a fresh large array that's exactly the same. But probably negligible in the grand scheme of things, within GLASS at least (the code predates use in GLASS).