jax

9923cefe - Merge branch 'vmap-improvements' of github.com:google/jax into vmap-improvements