xla

0d8a4c78 - cache the output sharding for dynamo