def _expand_dims_nonnegative_axis(axis, rank):
  """Get the nonnegative axis according to the rules of tf.expand_dims."""
  # Implementation note: equivalent to get_positive_axis(axis, rank + 1)
  if axis < 0:
    new_axis = (1 + rank) + axis
    if new_axis < 0:
      # Note: this is unreachable in the current code.
      raise ValueError("Axis out of range: " + str(axis))
    return new_axis
  elif axis > rank:
    # Note: this is unreachable in the current code.
    raise ValueError("Axis larger than rank: " + str(axis) + " > " + str(rank))
  return axis