jax
0e076992 - Simplify psum_invariant's abstract_eval rule by disallowing positional axes.

Commit

5 days ago

Simplify psum_invariant's abstract_eval rule by disallowing positional axes. PiperOrigin-RevId: 860159086

References

#26604 - [pallas:triton] Fix atomic min/max lowering for uint and float types

#34450 - [ROCm] Upstream ROCm CI Nightly Wheel Testing

#34607 - Simplify psum_invariant's abstract_eval rule by disallowing positional axes.

#34615 - [ROCm] Skip test_batch_axis_sharding_jvp on ROCm

#34675 - [ROCm] Update Skip Reason Output for run_on_devices decorator

Author

yashk2810

yashk2810

Committer

Google-ML-Automation

Google-ML-Automation

Parents

Loading