xla

a537020f - Apply sharding to all compatible input tensors (#4919)