code: def shift_right(x, axis=1):
      """Shift the input to the right for a given axis."""
      pad_widths = [(0, 0)] * len(x.shape)
      pad_widths[axis] = (1, 0)
      padded = jnp.pad(
          x,
          pad_widths,
          mode="constant",
          constant_values=x.dtype.type(0),
      )
      slicing = [slice(None)] * len(x.shape)
      slicing[axis] = slice(0, -1)
      return padded[tuple(slicing)]
  
imports: from typing import Any, Dict, Optional, Tuple
  from absl import logging
  from flax import linen
  from flax.training import common_utils
  from immutabledict import immutabledict
  import jax
  from jax import lax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.mbt import model
  from scenic.projects.mbt.model import temporal_encode
label_desc: Use image node_id 1 for calling a global function with example usage: shift_right(x, axis) and returns: padded