| 2133 | } |
| 2134 | |
| 2135 | XlaOp XlaBuilder::CrossReplicaSum( |
| 2136 | XlaOp operand, absl::Span<const ReplicaGroup> replica_groups) { |
| 2137 | return ReportErrorOrReturn([&]() -> StatusOr<XlaOp> { |
| 2138 | TF_ASSIGN_OR_RETURN(const Shape* shape, GetShapePtr(operand)); |
| 2139 | const Shape* element_shape; |
| 2140 | if (shape->IsTuple()) { |
| 2141 | if (shape->tuple_shapes_size() == 0) { |
| 2142 | return Unimplemented( |
| 2143 | "0 element tuple CrossReplicaSum is not supported"); |
| 2144 | } |
| 2145 | element_shape = &shape->tuple_shapes(0); |
| 2146 | } else { |
| 2147 | element_shape = shape; |
| 2148 | } |
| 2149 | const Shape scalar_shape = |
| 2150 | ShapeUtil::MakeShape(element_shape->element_type(), {}); |
| 2151 | auto b = CreateSubBuilder("sum"); |
| 2152 | auto x = b->Parameter(/*parameter_number=*/0, scalar_shape, "x"); |
| 2153 | auto y = b->Parameter(/*parameter_number=*/1, scalar_shape, "y"); |
| 2154 | if (scalar_shape.element_type() == PRED) { |
| 2155 | Or(x, y); |
| 2156 | } else { |
| 2157 | Add(x, y); |
| 2158 | } |
| 2159 | TF_ASSIGN_OR_RETURN(auto computation, b->Build()); |
| 2160 | return AllReduce(operand, computation, replica_groups, |
| 2161 | /*channel_id=*/absl::nullopt); |
| 2162 | }); |
| 2163 | } |
| 2164 | |
| 2165 | XlaOp XlaBuilder::AllReduce(XlaOp operand, const XlaComputation& computation, |
| 2166 | absl::Span<const ReplicaGroup> replica_groups, |
no test coverage detected