tagordon / spotlc Goto Github PK

I see you assert float64:

Is that required?

Just a quick note about using array.at[idx].set, like here:

There are some caveats about how this approach to ~in-place array updates works in jax to be aware of. I don't think this example will be a problem as long as we jit the methods that use it: https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html#jax.numpy.ndarray.at