I currently use JAX, but am interested in experimenting with Mojo. Is there a guide that helps JAX users transition to using Mojo? For example, it would be helpful to see how JAX's
vmap
vmap
/
shmap
shmap
operations could be expressed using Mojo syntax.
(I realize that Mojo is still under development, but given the rising popularity of JAX it might not be a bad idea to have some documentation helping people make the switch )