jax
c99d000a - Expose all_gather_invariant via `jax.lax.all_gather(..., to='invarying')` so that users have a way to go to invarying when the outside shard_map mesh context is Auto since 'reduced' doesn't work in Auto mode.

Commit

39 days ago

Expose all_gather_invariant via `jax.lax.all_gather(..., to='invarying')` so that users have a way to go to invarying when the outside shard_map mesh context is Auto since 'reduced' doesn't work in Auto mode. PiperOrigin-RevId: 891961880

References

#36352 - Expose all_gather_invariant via `jax.lax.all_gather(..., to='invarying')` so that users have a way to go to invarying when the outside shard_map mesh context is Auto since 'reduced' doesn't work in Auto mode.

Author

yashk2810

Committer

Google-ML-Automation

Parents

fdc12653

jax c99d000a - Expose all_gather_invariant via `jax.lax.all_gather(..., to='invarying')` so that users have a way to go to invarying when the outside shard_map mesh context is Auto since 'reduced' doesn't work in Auto mode.

jax
c99d000a - Expose all_gather_invariant via `jax.lax.all_gather(..., to='invarying')` so that users have a way to go to invarying when the outside shard_map mesh context is Auto since 'reduced' doesn't work in Auto mode.