def shift_right(x, axis=1):
  """Shift the input to the right by padding in the end on axis 1."""
  pad_widths = [(0, 0)] * len(x.shape)
  pad_widths[axis] = (0, 1)
  padded = jnp.pad(
      x[:, 1:], pad_widths, mode='constant', constant_values=x.dtype.type(0))
  return padded
Shift the input to the right by padding in the end on axis 1.
language/gscan/xattn_model/model/model_utils.py