xla

fc03aef3 - Update test_xla_sharding.py to take into account the number of devices (#8001)