def _rel_shift(F, x):
    """Perform relative shift operation following Dai et al. (2019) Appendix B.

    Unlike Dai et al.'s implementation, the relative shift is performed on the
    last two dimensions of the ndarray x. Further, we follow Yang et al. (2019)
    in not performing zero-padding but expecting the input array to be one
    element longer along the to be shifted dimension. For example, for
    TransformerXL, the pos_seq should be `arange(start=klen, stop=-qlen,
    step=-1)` instead of `arange(start=klen - 1, stop=-qlen, step=-1)`.

    Assumes len(x.shape) == 3 (could be generalized once F.swapaxes supports
    negative indices)

    """
    x_ = x
    # Reshape to x.shape[:-2] + [x.shape[-1] + 1, x.shape[-2]]
    x_ = F.reshape_like(x_, F.swapaxes(x_, 1, 2))
    # Remove padded elements
    x_ = F.slice_axis(x_, axis=-2, begin=1, end=None)
    # Reshape back to original shape
    x = F.reshape_like(x_, F.swapaxes(x_, 1, 2))
    return x