static */
| 2038 | } |
| 2039 | |
| 2040 | /* static */ StatusOr<Shape> ShapeInference::InferAllReduceShape( |
| 2041 | absl::Span<const Shape* const> operand_shapes) { |
| 2042 | for (const Shape* operand_shape : operand_shapes) { |
| 2043 | TF_RETURN_IF_ERROR( |
| 2044 | ExpectArray(*operand_shape, "operand of cross replica sum")); |
| 2045 | } |
| 2046 | if (operand_shapes.size() == 1) { |
| 2047 | return *operand_shapes[0]; |
| 2048 | } |
| 2049 | std::vector<Shape> operand_shape_values; |
| 2050 | for (const Shape* operand_shape : operand_shapes) { |
| 2051 | operand_shape_values.push_back(*operand_shape); |
| 2052 | } |
| 2053 | return ShapeUtil::MakeTupleShape(operand_shape_values); |
| 2054 | } |
| 2055 | |
| 2056 | /* static */ StatusOr<Shape> ShapeInference::InferAllToAllShape( |
| 2057 | const Shape& shape, int64 split_dimension, int64 concat_dimension, |
nothing calls this directly
no test coverage detected