commit 18a7a327f37b58b250d89190fe3cda9fa9372fab
author Swapnil Parekh <swapnilbp100@gmail.com> Wed May 20 19:13:15 2020 +0530
committer GitHub <noreply@github.com> Wed May 20 19:13:15 2020 +0530
tree 306fe939989d7bb6a814d23fc554216fa95f9627
parent cdd4e5e918a489778805ecb6fd8a9767e964d110

Update tpu_strategy.py

rank[DEP] -> ndim

tensorflow/python/distribute/tpu_strategy.py

diff --git a/tensorflow/python/distribute/tpu_strategy.py b/tensorflow/python/distribute/tpu_strategy.py
index b574c52..a8ffa61 100644
--- a/tensorflow/python/distribute/tpu_strategy.py
+++ b/tensorflow/python/distribute/tpu_strategy.py
@@ -897,7 +897,7 @@
           if tensor_util.is_tensor(input_tensor):
             rank = input_tensor.get_shape().rank
           else:
-            rank = np.rank(input_tensor)
+            rank = np.ndim(input_tensor)
           maximum_shape = tensor_shape.TensorShape([None] * rank)
           maximum_shapes.append(maximum_shape)
         maximum_shapes = nest.pack_sequence_as(replicate_inputs[0],