jax

dadc68b6 - add experimental lax.optimization_barrier autodiff rules