jax

77e4d8b3 - Updates onp -> np in random, loops, jet and in the tests of stax and optix (#3182)