jax
48fcf028 - Rename `psum2` to `psum_invariant` and put it in `lax_parallel`. We shouldn't expose this to public API and have users use `psum` instead which will dispatch to `psum_invariant` when `check_rep=True`.

Commit

287 days ago

Rename `psum2` to `psum_invariant` and put it in `lax_parallel`. We shouldn't expose this to public API and have users use `psum` instead which will dispatch to `psum_invariant` when `check_rep=True`. PiperOrigin-RevId: 745352875

Author

yashk2810

Committer

charleshofer

Parents

ef84e9d3

jax 48fcf028 - Rename `psum2` to `psum_invariant` and put it in `lax_parallel`. We shouldn't expose this to public API and have users use `psum` instead which will dispatch to `psum_invariant` when `check_rep=True`.

jax
48fcf028 - Rename `psum2` to `psum_invariant` and put it in `lax_parallel`. We shouldn't expose this to public API and have users use `psum` instead which will dispatch to `psum_invariant` when `check_rep=True`.