jax
64c0f62e - Sort manual axes when lowering `jax.shard_map` to `sdy.manual_computation`, which ensures the determinism in the generated `sdy.manual_computation`.

Commit

1 year ago

Sort manual axes when lowering `jax.shard_map` to `sdy.manual_computation`, which ensures the determinism in the generated `sdy.manual_computation`. PiperOrigin-RevId: 712973327

References

#25742 - Sort manual axes when lowering `jax.shard_map` to `sdy.manual_computation`, which ensures the determinism in the generated `sdy.manual_computation`.

#196 - CI: 01/08/25 upstream sync

Author

ZixuanJiang

Committer

Google-ML-Automation

Parents

00c363e1

jax 64c0f62e - Sort manual axes when lowering `jax.shard_map` to `sdy.manual_computation`, which ensures the determinism in the generated `sdy.manual_computation`.

jax
64c0f62e - Sort manual axes when lowering `jax.shard_map` to `sdy.manual_computation`, which ensures the determinism in the generated `sdy.manual_computation`.