def ShiftRightCls(cls_id):
  """Shifts right and insert cls.

  Args:
    cls_id: id of the cls token in embedding dictionary.  Returns a layer that
      shifts input tokens to the right by one and inserts an cls token to the
      beginning like in BERT paper.

  Returns:
    layer shifting to right and inserting cls.
  """

  def shift_right(x):
    pad_widths = [(0, 0)] * len(x.shape)
    pad_widths[1] = (1, 0)
    padded = jnp.pad(
        x, pad_widths, mode='constant', constant_values=x.dtype.type(cls_id))
    return padded[:, :-1]

  return cb.Fn('ShiftRightCls()', shift_right)
Shifts right and insert cls.

  Args:
    cls_id: id of the cls token in embedding dictionary.  Returns a layer that
      shifts input tokens to the right by one and inserts an cls token to the
      beginning like in BERT paper.

  Returns:
    layer shifting to right and inserting cls.
  
trax/layers/research/rel_attention.py