code: 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)
  
imports: from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import combinators
  from trax.layers import core
  from trax.layers import initializers
  from trax.layers.assert_shape import assert_shape
  from trax.layers.attention import MergeHeads
  from trax.layers.attention import SplitIntoHeads
label_desc: Use image node_id 10 for calling a global function with example usage: ShiftRightCls(cls_id) and returns: cb, padded