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.zeros_like(input_ids)
      shifted_input_ids = shifted_input_ids.at[:, 1:].set(
          input_ids[:, :-1]
      )
      shifted_input_ids = shifted_input_ids.at[:, 0].set(
          decoder_start_token_id
      )
  
      shifted_input_ids = jnp.where(
          shifted_input_ids == -100, pad_token_id, shifted_input_ids
      )
      return shifted_input_ids
  
file_path: transformers/src/transformers/models/mt5/modeling_flax_mt5.py
imports: import jax.numpy
  from ...utils import logging
  from ..t5.modeling_flax_t5 import FlaxT5EncoderModel, FlaxT5ForConditionalGeneration, FlaxT5Model
  from .configuration_mt5 import MT5Config
label_desc: Use image node_id 1 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