def create_unbroadcast_axis(shape, broadcast_shape):
  """Creates the reduction axis for unbroadcasting.

  Args:
    shape: A list. The shape after the broadcast operation.
    broadcast_shape: A list. The original shape the array being unbroadcast
      had.
  Returns:
    A list. The axes along which the array needs to be reduced. These axes will
    be distributed evenly into the original shape.
  """
  return tuple(
      -(1 + i)
      for i in range(len(broadcast_shape))
      if i >= len(shape) or broadcast_shape[-(1 + i)] > shape[-(1 + i)])