message: Fix a bug that if all the input shapes to TPUStrategy.experimental_run_v2 is static, it always uses the shape from first replica to pad. But it may fail if the first replica shape is not the maximum shape.
diff: @@ -675,9 +675,10 @@ class TPUExtended(distribute_lib.StrategyExtendedV1):
  flattened_list = nest.flatten(replicate_inputs[0])
  for input_tensor in flattened_list:
  if tensor_util.is_tensor(input_tensor):
  - maximum_shape = input_tensor.get_shape()
  + rank = input_tensor.get_shape().rank
  else:
  - maximum_shape = tensor_shape.TensorShape(np.shape(input_tensor))
  + rank = np.rank(input_tensor)
  + maximum_shape = tensor_shape.TensorShape([None] * rank)
  maximum_shapes.append(maximum_shape)
  maximum_shapes = nest.pack_sequence_as(replicate_inputs[0],
  maximum_shapes)
  