Run persistent cache test on TPU CI

[jonb377]

Persistent cache test only ran on GPU CI. This change enables it on TPU CI as well.

[JackCaoG]

Feel free to merge after TPU CI passed

[alanwaketan]

LGTM.

[jonb377]

@JackCaoG Looks like CI failed with timeout, was this the same issue you hit?

Step #4 - "run_e2e_tests": + python3 test/test_persistent_cache.py
Step #4 - "run_e2e_tests": Running tests under Python 3.10.14: /usr/local/bin/python3
Step #4 - "run_e2e_tests": [ RUN      ] PersistentCacheTest.test_persistent_cache_mp
Step #4 - "run_e2e_tests": [  FAILED  ] PersistentCacheTest.test_persistent_cache_mp
Step #4 - "run_e2e_tests": [ RUN      ] PersistentCacheTest.test_persistent_cache_single_device
Step #4 - "run_e2e_tests": ++ kubectl get pod/xla-test-job-zqtd5 -o 'jsonpath={.status.containerStatuses[?(@.name=="xla-test")].state.terminated.exitCode}'
Step #4 - "run_e2e_tests": + exit 137

I can repro locally with the XLA_USE_SPMD env var set, but it passes without that env var.