xla

1e265c4a - move prepareoutput sharding before the lock region