xla

8efb0b53 - [SPMD] Properly skip tests on TPU V2

[SPMD] Properly skip tests on TPU V2 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