-
Notifications
You must be signed in to change notification settings - Fork 3.2k
Open
Labels
Description
Named collectives have come up before in documentation for pmap/shmap, but usually links to https://jax.readthedocs.io/en/latest/jax.lax.html#parallel-operators (other relevant issue: #13794). However, these also work for things such as vmap and I think warrant more documentation, since I discovered them via: #22974 (reply in thread). My proposal is to have a list of all collectives that work under vmap (somewhere) in the docs.