jax

0293ecbb - Add support for vmap of scatter where indices but not updates are batched.