class_docstr: An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
  models.
class_docstr_tok: ["An","abstract","class","to","handle","weights","initialization","and","a","simple","interface","for","downloading","and","loading","pretrained","models","."]
code: def __call__(
      self,
      input_ids,
      attention_mask=None,
      position_ids=None,
      params: dict = None,
      past_key_values: dict = None,
      dropout_rng: jax.random.PRNGKey = None,
      train: bool = False,
      output_attentions: Optional[bool] = None,
      output_hidden_states: Optional[bool] = None,
      return_dict: Optional[bool] = None,
  ):
      output_attentions = (
          output_attentions
          if output_attentions is not None
          else self.config.output_attentions
      )
      output_hidden_states = (
          output_hidden_states
          if output_hidden_states is not None
          else self.config.output_hidden_states
      )
      return_dict = (
          return_dict
          if return_dict is not None
          else self.config.return_dict
      )
  
      batch_size, sequence_length = input_ids.shape
  
      if position_ids is None:
          if past_key_values is not None:
              raise ValueError(
                  "Make sure to provide `position_ids` when passing `past_key_values`."
              )
  
          position_ids = jnp.broadcast_to(
              jnp.arange(sequence_length)[None, :],
              (batch_size, sequence_length),
          )
  
      if attention_mask is None:
          attention_mask = jnp.ones((batch_size, sequence_length))
  
      # Handle any PRNG if needed
      rngs = {}
      if dropout_rng is not None:
          rngs["dropout"] = dropout_rng
  
      inputs = {"params": params or self.params}
  
      # if past_key_values are passed then cache is already initialized a private flag init_cache has to be passed down to ensure cache is used. It has to be made sure that cache is marked as mutable so that it can be changed by FlaxLlamaAttention module
      if past_key_values:
          inputs["cache"] = past_key_values
          mutable = ["cache"]
      else:
          mutable = False
  
      outputs = self.module.apply(
          inputs,
          jnp.array(input_ids, dtype="i4"),
          jnp.array(attention_mask, dtype="i4"),
          jnp.array(position_ids, dtype="i4"),
          not train,
          False,
          output_attentions,
          output_hidden_states,
          return_dict,
          rngs=rngs,
          mutable=mutable,
      )
  
      # add updated cache to model output
      if past_key_values is not None and return_dict:
          outputs, past_key_values = outputs
          outputs["past_key_values"] = unfreeze(
              past_key_values["cache"]
          )
          return outputs
      elif past_key_values is not None and not return_dict:
          outputs, past_key_values = outputs
          outputs = (
              outputs[:1]
              + (unfreeze(past_key_values["cache"]),)
              + outputs[1:]
          )
  
      return outputs
  
imports: from functools import partial
  from typing import Optional, Tuple
  import flax.linen
  import jax
  import jax.numpy
  import 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 ...modeling_flax_outputs import FlaxBaseModelOutput, FlaxCausalLMOutput
  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_llama import LlamaConfig
label_desc: Use image node_id 4 for calling the FlaxLlamaPreTrainedModel obj's underlying member method code with example usage: obj.__call__(input_ids, attention_mask, position_ids, params, past_key_values, dropout_rng, train, output_attentions, output_hidden_states, return_dict) and returns: outputs, outputs