Comment by breakds

For the second problem, if I understand it correctly, you might want to try `vmap` from jax.