Comment by semiinfinitely

all issues raised are addressed by jax.vmap