[XLA] Add the xla interface for AllToAll.

PiperOrigin-RevId: 207971529

1 files changed, 12 insertions, 1 deletions

diff --git a/tensorflow/compiler/xla/service/shape_inference.h b/tensorflow/compiler/xla/service/shape_inference.h
index 33da323b3d..c185b0a1bd 100644
--- a/tensorflow/compiler/xla/service/shape_inference.h
+++ b/tensorflow/compiler/xla/service/shape_inference.h
@@ -119,11 +119,22 @@ class ShapeInference {
       const Shape& in, FftType fft_type,
       tensorflow::gtl::ArraySlice<int64> fft_length);
 
-  // Infers the shape produced a cross replica sum with the given operand
+  // Infers the shape produced by a cross replica sum with the given operand
   // shapes.
   static StatusOr<Shape> InferCrossReplicaSumShape(
       tensorflow::gtl::ArraySlice<const Shape*> operand_shapes);
 
+  // Infers final shape of an Alltoall operation that is created by the xla
+  // builder.
+  static StatusOr<Shape> InferAllToAllShape(const Shape& shape,
+                                            int64 split_dimension,
+                                            int64 concat_dimension,
+                                            int64 split_count);
+
+  // Infers the shape of an HLO all-to-all instruction.
+  static StatusOr<Shape> InferAllToAllTupleShape(
+      tensorflow::gtl::ArraySlice<const Shape*> operand_shapes);
+
   // Infers the shape produced by applying the given reduction computation
   // shape to the given input operand shape.
   //