xla

6d88c089 - Use shard_as in scan to ensure that inputs and their gradients have the same sharding (#8879)