jax
8a73c61b - [Mosaic TPU] Support reshape which unfolds the minormost dim into two dims when the minormost dim is not divisible by 128, also speed up the reshape for already supported cases where the minormost dim is a multiple of 128.

Commit

49 days ago

[Mosaic TPU] Support reshape which unfolds the minormost dim into two dims when the minormost dim is not divisible by 128, also speed up the reshape for already supported cases where the minormost dim is a multiple of 128. PiperOrigin-RevId: 878092970

References

#35570 - [Mosaic TPU] Support reshape which unfolds the minormost dim into two dims when the minormost dim is not divisible by 128, also speed up the reshape for already supported cases where the minormost dim is a multiple of 128.

Author

yueshengys

Committer

Google-ML-Automation

Parents

d777b52f

jax 8a73c61b - [Mosaic TPU] Support reshape which unfolds the minormost dim into two dims when the minormost dim is not divisible by 128, also speed up the reshape for already supported cases where the minormost dim is a multiple of 128.

jax
8a73c61b - [Mosaic TPU] Support reshape which unfolds the minormost dim into two dims when the minormost dim is not divisible by 128, also speed up the reshape for already supported cases where the minormost dim is a multiple of 128.