Comment by _hark

jax.vmap() is all you need?