Fix layout assignment for cross module all reduce
Previously we could have ended up with the different HLOs being assigned different layouts what made lowering impossible. This change enforces a consistent layout between the communicating nodes the same way it is done for send&recv pairs. PiperOrigin-RevId: 215359420
Loading
Please sign in to comment