code: def assert_rank(tensor, expected_rank, name=None):
      """Raises an exception if the tensor rank is not of the expected rank.
  
      Args:
        tensor: A tf.Tensor to check the rank of.
        expected_rank: Python integer or list of integers, expected rank.
        name: Optional name of the tensor for the error message.
  
      Raises:
        ValueError: If the expected shape doesn't match the actual shape.
      """
      expected_rank_dict = {}
      if isinstance(expected_rank, six.integer_types):
          expected_rank_dict[expected_rank] = True
      else:
          for x in expected_rank:
              expected_rank_dict[x] = True
  
      actual_rank = tensor.shape.ndims
      if actual_rank not in expected_rank_dict:
          raise ValueError(
              "For the tensor `%s`, the actual tensor rank `%d` (shape = %s) is not "
              "equal to the expected tensor rank `%s`"
              % (
                  name,
                  actual_rank,
                  str(tensor.shape),
                  str(expected_rank),
              )
          )
  
file_path: autodist/examples/benchmark/utils/tf_utils.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import six
  import tensorflow
  from utils import activations
label_desc: Use image node_id 6 for calling a global function with example usage: assert_rank(tensor, expected_rank, name) without return types