Merge pull request #52863 from tensorflow/fix-build-3-on-r2.5

Fix build on `tpu_cross_replica_ops.cc`
tensorflow · Oct 29, 2021 · 2620d2c · 2620d2c
2 parents a7ef5db + 9bd7e1e
commit 2620d2c
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/tensorflow/core/ops/tpu_cross_replica_ops.cc b/tensorflow/core/ops/tpu_cross_replica_ops.cc
@@ -33,10 +33,14 @@ REGISTER_OP("AllToAll")
     .SetShapeFn([](InferenceContext* c) {
       ShapeHandle input = c->input(0);
       ShapeHandle group_assignment = c->input(1);
-      if (!c->RankKnown(input)) {
-        c->set_output(0, c->UnknownShape());
-        return Status::OK();
+
+      int64 rank;
+      if (c->RankKnown(input)) {
+        rank = c->Rank(input);
+      } else {
+        return errors::InvalidArgument("input's rank is unknown.");
       }
+
       int concat_dimension;
       int split_dimension;
       int split_count;