xla

aab2b80f - Cache the output sharding spec for dynamo