jax
48140bcb - Add basic support of `unreduced` to sharding-in-types! We cannot lower it right now, but it atleast shows up in types.

Commit

264 days ago

Add basic support of `unreduced` to sharding-in-types! We cannot lower it right now, but it atleast shows up in types. The API to specify unreduced is via `PartitionSpec`. For example: `PartitionSpec('x', 'y', None, unreduced='z')` or `PartitionSpec('x', unreduced=('y', 'z'))`. In types/jaxpr, unreduced will show up as: `f32[8@x,2]{U:y}` But we support unreduced only in dot_general and nary ops (add, mul, etc) as of this change: (the support will be expanded in following changes) * **dot general** only allows unreduced when contracting dims are sharded. And the unreduced axes specified by the user needs to match the sharding of the contracting dims. In all other cases, an error is raised. An example of how unreduced can be specified: `jnp.einsum('xy,yz->xz', x, y, out_sharding=P('x', unreduced='y'))` * **nary ops** can propagate unreduced (add, mul, etc). If all ops aren't unreduced across the same mesh axes, an error is raised. PiperOrigin-RevId: 756063074

References

#28597 - Add basic support of `unreduced` to sharding-in-types! We cannot lower it right now, but it atleast shows up in types.

#477 - CI: 06/20/25 upstream sync

#31381 - Remove old ROCm build code

#31720 - Fix ann_test.py numerical bug in target reshape

#31768 - [ROCm] Support lowering through PJRT_Triton_Extension

#32115 - Relax version requirements for ROCm Jax Plugin wheels

#33157 - [ROCm] Resolve undefined behavior in bitshift unit test

#579 - Create rocm-test-requirements.txt

#581 - Fix/pallas tests shared memory

#584 - Use plain bazel to test jax, use hermetic rocm dependency

#585 - update a test for checking zero ROCm GPU event