import torch


def flip(x, dim):
    """Reverse tensor order aling given dimension.
    Drop-in replacement for tf.reverse(x, axis)

    Parameters
    ----------
    x : torch.Tensor
        tensor to flip
    dim : int
        dimension along which to flip

    Returns
    -------
    torch.Tensor
        flipped tensor

    References
    ----------

    [1] : https://github.com/pytorch/pytorch/issues/229
    """
    indices = [slice(None)] * x.dim()
    indices[dim] = torch.arange(
        x.size(dim) - 1, -1, -1, dtype=torch.long, device=x.device
    )
    return x[tuple(indices)]