[Pallas] Set the major-minor layout for inputs and outputs (#6826)
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
