From 9bd7e1e87153a69353ba652efce7e81ca0e34b75 Mon Sep 17 00:00:00 2001
From: Mihai Maruseac <mihaimaruseac@google.com>
Date: Fri, 29 Oct 2021 09:55:27 -0700
Subject: [PATCH] Fix build on `tpu_cross_replica_ops.cc`

---
 tensorflow/core/ops/tpu_cross_replica_ops.cc | 10 +++++++---
 1 file changed, 7 insertions(+), 3 deletions(-)

diff --git a/tensorflow/core/ops/tpu_cross_replica_ops.cc b/tensorflow/core/ops/tpu_cross_replica_ops.cc
index c2f4ab73986035..404c3e59eff941 100644
--- 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;