That's a great answer, thank you :) Stuff like vmap and grad and pmap and all the rest have been a huge boon in simplifying some of my work, so I'm glad to see

That's a great answer, thank you :) Stuff like vmap and grad and pmap and all the rest have been a huge boon in simplifying some of my work, so I'm glad to see it's expanding into pytorch!