def _validate_init_input_and_get_axis(reference_batch, axis):
  """Validate input and return the used axis value."""
  if reference_batch.shape.ndims is None:
    raise ValueError('`reference_batch` has unknown dimensions.')

  ndims = reference_batch.shape.ndims
  if axis < 0:
    used_axis = ndims + axis
  else:
    used_axis = axis
  if used_axis < 0 or used_axis >= ndims:
    raise ValueError('Value of `axis` argument ' + str(used_axis) +
                     ' is out of range for input with rank ' + str(ndims))
  return used_axis