def make_non_negative_axis(axis, rank):
  """Make (possibly negatively indexed) `axis` argument non-negative."""
  axis = tf.convert_to_tensor(value=axis, name="axis")
  rank = tf.convert_to_tensor(value=rank, name="rank")
  axis_ = tf.get_static_value(axis)
  rank_ = tf.get_static_value(rank)

  # Static case.
  if axis_ is not None and rank_ is not None:
    is_scalar = axis_.ndim == 0
    if is_scalar:
      axis_ = [axis_]
    positive_axis = []
    for a_ in axis_:
      if a_ < 0:
        positive_axis.append(rank_ + a_)
      else:
        positive_axis.append(a_)
    if is_scalar:
      positive_axis = positive_axis[0]
    return tf.convert_to_tensor(value=positive_axis, dtype=axis.dtype)

  # Dynamic case.
  # Unfortunately static values are lost by this tf.where.
  return tf.where(axis < 0, rank + axis, axis)
Make (possibly negatively indexed) `axis` argument non-negative.