code: def shift_tokens_right(
      input_ids: jnp.ndarray,
      pad_token_id: int,
      decoder_start_token_id: int,
  ) -> jnp.ndarray:
      """
      Shift input ids one token to the right.
      """
      shifted_input_ids = jnp.roll(input_ids, 1, axis=-1)
      shifted_input_ids = shifted_input_ids.at[(..., 0)].set(
          decoder_start_token_id
      )
      # replace possible -100 values in labels by `pad_token_id`
      shifted_input_ids = jnp.where(
          shifted_input_ids == -100, pad_token_id, shifted_input_ids
      )
  
      return shifted_input_ids
  
imports: import math
  import random
  from functools import partial
  from typing import Optional, Tuple
  import numpy
  import flax.linen
  import jax
  import jax.numpy
  from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
  from flax.linen import combine_masks, make_causal_mask
  from flax.linen.attention import dot_product_attention_weights
  from flax.traverse_util import flatten_dict, unflatten_dict
  from jax import lax
  from jax.random import PRNGKey
  from ...modeling_flax_outputs import FlaxBaseModelOutputWithPastAndCrossAttentions, FlaxCausalLMOutputWithCrossAttentions
  from ...modeling_flax_utils import ACT2FN, FlaxPreTrainedModel, append_call_sample_docstring
  from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging
  from .configuration_xglm import XGLMConfig
label_desc: Use image node_id 2 for calling a global function with example usage: shift_tokens_right(input_ids, pad_token_id, decoder_start_token_id) and returns: shifted_input_ids