jax
4c671ca7 - [Pallas MGPU] Adding support for squeezed block dims in the pipeline BlockSpecs. They can be identified with a `None` or `pl.Squeezed`.

Commit

62 days ago

[Pallas MGPU] Adding support for squeezed block dims in the pipeline BlockSpecs. They can be identified with a `None` or `pl.Squeezed`. PiperOrigin-RevId: 845782837

References

#33713 - [Pallas MGPU] Adding support for squeezed block dims in the pipeline BlockSpecs. They can be identified with a `None` or `pl.Squeezed`.

Author

Rifur13

Committer

Google-ML-Automation

Parents

44f67b03

jax 4c671ca7 - [Pallas MGPU] Adding support for squeezed block dims in the pipeline BlockSpecs. They can be identified with a `None` or `pl.Squeezed`.

jax
4c671ca7 - [Pallas MGPU] Adding support for squeezed block dims in the pipeline BlockSpecs. They can be identified with a `None` or `pl.Squeezed`.