jax

0a98102a - Merge pull request #37928 from lanarkite99:resume/jax-bcoo-vmap-grad