xla

c3a59095 - Revert "Remove reshape lowering (#6230)" (#6280)

Revert "Remove reshape lowering (#6230)" (#6280) This reverts commit 0c81365. It causes OOMs for some training workloads on GPUs, and slows down others.