xla

bf997819 - Fix the input sharding for dynamo (#5469)