Getting Started
Reference Documentation
Advanced Topics
Examples
API Documentation
list_pmap()
pmaps a function over similar pytrees.
fun – Function accepting a single pytree as first argument.
Returns a pmapped-function accepting multiple pytree args with similar tree- structure.