xla

Use shard_as in scan to ensure that inputs and their gradients have the same sharding

#8879

Merged