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.
      """
      if name is None:
          name = tensor.name
  
      expected_rank_dict = {}
      if isinstance(expected_rank, int):
          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:
          scope_name = tf.compat.v1.get_variable_scope().name
          raise ValueError(
              "For the tensor `%s` in scope `%s`, the actual rank "
              "`%d` (shape = %s) is not equal to the expected rank `%s`"
              % (
                  name,
                  scope_name,
                  actual_rank,
                  str(tensor.shape),
                  str(expected_rank),
              )
          )
  
imports: import copy
  import json
  import math
  import numpy
  import tensorflow
  import functools
  from transformers.activations_tf import gelu
  from finetune.optimizers.recompute_grads import recompute_grad
  from finetune.nn.auxiliary import embed_position
label_desc: Use image node_id 19 for calling a global function with example usage: assert_rank(tensor, expected_rank, name) without return types