[XLA::SPMD] Propagating shardings between operands of element-wise operations.
Let us use C = foo(A, B) as an example. We have two key steps in sharding propagation.
Forward propagation. Infer C's sharding from A and B.
Backward propagation. Infer A's or B's sharding from C.
Ideally, we do not need to explicitly propagate shardings between operands since we have the paths A -> C -> B and B -> C -> A. However, this path may be unavailable if C has a pre-defined sharding, which impedes the forward propagation. To resolve this issue, we may add the propagation A -> B and B -> A directly.
It was already considered for several operations, such as dot and reduce. This cl support it for element-wise operations.
[XLA::SPMD] Propagating shardings between operands of element-wise operations.
Let us use
C = foo(A, B)as an example. We have two key steps in sharding propagation.Ideally, we do not need to explicitly propagate shardings between operands since we have the paths
A -> C -> BandB -> C -> A. However, this path may be unavailable if C has a pre-defined sharding, which impedes the forward propagation. To resolve this issue, we may add the propagationA -> BandB -> Adirectly.It was already considered for several operations, such as dot and reduce. This cl support it for element-wise operations.