xla

364819a9 - [SPMD] Properly skip tests on TPU V2 (#5479)

[SPMD] Properly skip tests on TPU V2 (#5479) Summary: Some of the tests only fail on TPU V2 but were skipped for all TPUs. Let's fix that. Test Plan: PJRT_DEVICE=TPU USE_XLA_SPMD=1 python test/spmd/test_xla_sharding.py