Searched refs:InferAllReduceShape (Results 1 – 4 of 4) sorted by relevance
128 static StatusOr<Shape> InferAllReduceShape(
192 return CheckShape(crs, ShapeInference::InferAllReduceShape(operand_shapes)); in HandleAllReduce()
1975 /* static */ StatusOr<Shape> ShapeInference::InferAllReduceShape( in InferAllReduceShape() function in xla::ShapeInference
2142 ShapeInference::InferAllReduceShape({&operand_shape})); in CrossReplicaSum()