Neat! It /should/ be pretty easy to convert this to Jax, as a big party of the Jax API is just reimplementing the Numpy API. The two changes I can imagine are:

Neat! It /should/ be pretty easy to convert this to Jax, as a big party of the Jax API is just reimplementing the Numpy API. The two changes I can imagine are: a) handling randomness requires more explicit rng estate handling, and b) wrapping the inner loop in a function so that it can be JIT-compiled. The advantage to doing this is that you can then run your code on a GPU, getting you most of the way to 'blazingly fast.'