xla

1c57aeae - Parallel tests (#4960) (#5030)

Parallel tests (#4960) (#5030) * Improve coverage config * Rearrange tests to speed them up. * Run TPU and GPU tests serially. * Automate sharding