jax
dc7d56e6 - [Mosaic TPU] Allow using replicated lane offset for intermediate safe offsets in mask relayout since relayout like {0,},(8,128) -> {0,},(4,128) has been just supported.

Commit

51 days ago

[Mosaic TPU] Allow using replicated lane offset for intermediate safe offsets in mask relayout since relayout like {0,*},(8,128) -> {0,*},(4,128) has been just supported. PiperOrigin-RevId: 826555654

References

#31381 - Remove old ROCm build code

#33157 - Resolve undefined behavior in bitshift unit test

#33186 - Make nvidia version data optional for ROCm builds

Author

yueshengys

Committer

Google-ML-Automation

Parents

7a007ea9

jax dc7d56e6 - [Mosaic TPU] Allow using replicated lane offset for intermediate safe offsets in mask relayout since relayout like {0,*},(8,128) -> {0,*},(4,128) has been just supported.

jax
dc7d56e6 - [Mosaic TPU] Allow using replicated lane offset for intermediate safe offsets in mask relayout since relayout like {0,},(8,128) -> {0,},(4,128) has been just supported.