def ShiftRight(n_positions=1, mode='train'):
  """Returns a layer that can insert padding to shift the input sequence.

  Args:
    n_positions: Number of positions to shift the input sequence rightward;
        initial positions freed by the shift get padded with zeros. Applies
        only if layer is created in a non-``'eval'`` mode.
    mode: One of ``'train'``, ``'eval'``, or ``'predict'``.
  """
  # TODO(jonni): Include pad arg, like PaddingMask, to allow non-default pads?
  def f(x):
    if mode == 'predict':
      return x
    padded = _zero_pad(x, (n_positions, 0), 1)
    return padded[:, :-n_positions]
  return Fn(f'ShiftRight({n_positions})', f)
Returns a layer that can insert padding to shift the input sequence.

  Args:
    n_positions: Number of positions to shift the input sequence rightward;
        initial positions freed by the shift get padded with zeros. Applies
        only if layer is created in a non-``'eval'`` mode.
    mode: One of ``'train'``, ``'eval'``, or ``'predict'``.
  
trax/layers/attention.py