reshape implementation

Commit

174 days ago

[Mosaic TPU] Refactor the reshape->store optimization to follow the load->reshape implementation It's the same trick, just applied in reverse. PiperOrigin-RevId: 829334513

References

#33085 - [Mosaic TPU] Refactor the reshape->store optimization to follow the load->reshape implementation

Author

apaszke

Committer

Google-ML-Automation

Parents

4cdb3cea

jax 062e6a41 - [Mosaic TPU] Refactor the reshape->store optimization to follow the load->reshape implementation

jax
062e6a41 - [Mosaic TPU] Refactor the reshape->store optimization to follow the load->reshape implementation