def get_dim_as_int(dim):
    """Utility to get v1 or v2 TensorShape dim as an int.

  Args:
    dim: The TensorShape dimension to get as an int

  Returns:
    None or an int.
  """
    try:
        return dim.value
    except AttributeError:
        return dim
