xla

1e4e5d23 - Update test_xla_sharding.py to take into account the number of devices