initial commit of Jax integration [WIP]

[iancharest]

Jax seems to offer quite substantial speed improvements on things like svd compared to numpy and scipy, even in the absence of GPU or TPU.

Let's see if this is the case. we will add jax to fracridge, and then test it with glm_denoisesingletrial

• add jax to requirements.
• try and catch jax, else use scipy

[iancharest]

ok, learning a bit about jax.
https://github.com/iancharest/fracridge/blob/46c3090ef47c241dc07a399f60d28dbb7ea9e0bf/fracridge/fracridge.py#L155
leads to:
TypeError: '<class 'jax.interpreters.xla.DeviceArray'>' object does not support item assignment. JAX arrays are immutable; perhaps you want jax.ops.index_update or jax.ops.index_add instead?