xla

3905372e - [Backport] Set the major-minor layout for inputs and outputs (#6876)

[Backport] Set the major-minor layout for inputs and outputs (#6876) Summary: Mosaic only accepts major-minor layout for both its inputs and outputs. So we need to enforce those layouts by setting the expected input&output shapes in xla::CustomCallWithLayout. After this change, XLA_TPU_LAYOUT is no longer needed. Test Plan: PJRT_DEVICE=TPU python test/test_pallas.py