JIT-compiled first- and second-order optimization algorithms in JAX.

[WIP] Nano Optax