The tf.TensorShape of a tensor.

    Args:
      tensor_name: string, the name of a tensor in the graph.

    Returns:
      a tf.TensorShape

def get_tensor_shape(self, tensor_name):
    """The tf.TensorShape of a tensor.

    Args:
      tensor_name: string, the name of a tensor in the graph.

    Returns:
      a tf.TensorShape
    """
    tensor = self._name_to_tensor(tensor_name)

    if isinstance(tensor, mtf.Tensor):
      return tf.TensorShape(tensor.shape.to_integer_list)
    else:  # tf.Tensor
      return tensor.shape