Will pmap be supported in functorh？

pytorch / functorch

functorch is JAX-like composable function transforms for PyTorch.

https://pytorch.org/functorch/

BSD 3-Clause "New" or "Revised" License

1.38k stars 102 forks source link

Will pmap be supported in functorh？ #1099

Open shixun404 opened 1 year ago

shixun404 commented 1 year ago

Greetings, I am very grateful that vmap is supported in functorch. Is there any plan to include support for pmap in the future? Thank you. Additionally, what are the ways that I can contribute to the development of this project?

jordiae commented 1 year ago

Any updates on this?

ablghtianyi commented 2 weeks ago

I am also curious about this. Having a pmap similar to the one in jax would be very handy for customizing functions.