jax
746e13ef - [export] Fix forward compatibility for exporting nr_devices

Commit

55 days ago

[export] Fix forward compatibility for exporting nr_devices In #33492 I have added a new 32-bit field to store the nr_devices, and renamed the old 16-bit field to nr_devices_short. However, that change mistakenly stops populating the 16-bit field, which means that when the export is read by old code, it will have a nr_devices == 0.

References

#31768 - [ROCm] Support lowering through PJRT_Triton_Extension

#33685 - [export] Fix forward compatibility for exporting nr_devices

#34135 - [ROCm] update to test if there are GPU events when doing profiling on…

#592 - Upstream ci download jax rocm wheels

#34304 - [ROCm] Add ROCm platform support for psend/precv collective operations

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

#34470 - [ROCm] Enabled testTridiagonal for ROCm devices.

#34474 - [ROCm] Added support for GESVDJ on ROCm devices

#34494 - [ROCm] Enabled ROCm devices to default to Jacobi SVD on smaller matrices.

#632 - Add script for running bazel tests on ROCm

#638 - Support wheels as build deps