import torch


def reverse(x, axis):
  """Reverse a tensor along the specified axes.

  Arguments:
      x: Tensor to reverse.
      axis: Integer or iterable of integers.
          Axes to reverse.

  Returns:
      A tensor.
  """
  if isinstance(axis, int):
    axis = [axis]
  return  torch.flip(x,dims=axis)