code: def shift_right(x, axis=1):
      """Shift the input to the right by padding and slicing on 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),
      )
      return lax.dynamic_slice_in_dim(
          padded, 0, padded.shape[axis] - 1, axis
      )
  
imports: from __future__ import annotations
  from typing import Callable, Any, Optional
  from flax import linen
  from flax import struct
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: shift_right(x, axis) and returns: lax