jax
02cad83c - [Mosaic:TPU] Have VectorLayout::join handle fully replicated layouts and fix bug in infer-vector-layout elementwise rule

Commit

95 days ago

[Mosaic:TPU] Have VectorLayout::join handle fully replicated layouts and fix bug in infer-vector-layout elementwise rule Consider the case where input layouts are #tpu.vpad<"32,{*,*},(8,128),-1"> #tpu.vpad<"32,{*,0},(8,128)"> in this case we would have previously assigned the second layout, but that is incorrect because the true 2nd minor dimension is not known to be replicated for the first input PiperOrigin-RevId: 810776770

References

#31381 - Remove old ROCm build code

#32115 - Relax version requirements for ROCm Jax Plugin wheels

#33157 - Resolve undefined behavior in bitshift unit test

#33186 - Make nvidia version data optional for ROCm builds

#579 - Create rocm-test-requirements.txt

#580 - GESVDJ support for ROCm GPUs in JAX

#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

Author