code: def shift_inputs(x, segment_ids=None, axis=1):
      """Shift inputs and replace EOS by 0 for packed inputs."""
      shifted = shift_right(x, axis=axis)
      # For packed targets, the first shifted token of a new sequence is made
      # 0, rather than being the EOS token for the last sequence.
      if segment_ids is not None:
          shifted *= segment_ids == shift_right(segment_ids, axis=axis)
      return shifted
  
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 2 for calling a global function with example usage: shift_inputs(x, segment_ids, axis) and returns: shifted