def _check_valid_index(index, element_name):
  """Verifies if a value with `element_name` is a valid index."""
  if isinstance(index, int):
    return True
  elif isinstance(index, tf.Tensor):
    if index.dtype != tf.int32 and index.dtype != tf.int64:
      raise TypeError(
          "Invalid tensor `{}` parameter. Valid tensor indices must have "
          "types tf.int32 or tf.int64, got {}."
          .format(element_name, index.dtype))
    if index.shape.as_list():
      raise TypeError(
          "Invalid tensor `{}` parameter. Valid tensor indices must be scalars "
          "with shape [], got{}"
          .format(element_name, index.shape.as_list()))
    return True
  else:
    raise TypeError(
        "Invalid `{}` parameter. Valid tensor indices must be integers "
        "or tensors, got {}."
        .format(element_name, type(index)))
Verifies if a value with `element_name` is a valid index.
graph_nets/utils_tf.py