xla

919f0aee - Apply sharding to all compatible input tensors (#4872)