jax
19c40e08 - [Mosaic TPU] Support general reshape patterns: `[src_major_dims, AxB] -> [dst_major_dims, A, B]`, where we can have arbitrary reshapes between `src_major_dims` and `dst_major_dim`.

Commit

78 days ago

[Mosaic TPU] Support general reshape patterns: `[src_major_dims, AxB] -> [dst_major_dims, A, B]`, where we can have arbitrary reshapes between `src_major_dims` and `dst_major_dim`. PiperOrigin-RevId: 882832471

References

#35753 - [Mosaic TPU] Support general reshape patterns: `[src_major_dims, AxB] -> [dst_major_dims, A, B]`, where we can have arbitrary reshapes between `src_major_dims` and `dst_major_dim`.

Author

yueshengys

Committer

Google-ML-Automation

Parents

20c8683e

jax 19c40e08 - [Mosaic TPU] Support general reshape patterns: `[src_major_dims, AxB] -> [dst_major_dims, A, B]`, where we can have arbitrary reshapes between `src_major_dims` and `dst_major_dim`.

jax
19c40e08 - [Mosaic TPU] Support general reshape patterns: `[src_major_dims, AxB] -> [dst_major_dims, A, B]`, where we can have arbitrary reshapes between `src_major_dims` and `dst_major_dim`.