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
Author
Parents
Loading