class_docstr: 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,
      past_key_values: dict = None,
      params: 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.use_return_dict
      )
  
      batch_size, sequence_length = input_ids.shape
  
      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 FlaxBloomAttention 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"),
          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: import math
  from functools import partial
  from typing import Optional, Tuple
  import flax.linen
  import jax
  import jax.numpy
  from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
  from flax.linen import combine_masks, dot_product_attention_weights, make_causal_mask
  from flax.linen.activation import tanh
  from flax.traverse_util import flatten_dict, unflatten_dict
  from jax import lax
  from ...modeling_flax_outputs import FlaxBaseModelOutput, FlaxBaseModelOutputWithPastAndCrossAttentions, FlaxCausalLMOutput
  from ...modeling_flax_utils import FlaxPreTrainedModel, append_call_sample_docstring
  from ...utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging
  from .configuration_bloom import BloomConfig
label_desc: Use image node_id 4 for calling the FlaxBloomPreTrainedModel obj's underlying member method code with example usage: obj.__call__(input_ids, attention_mask, past_key_values, params, dropout_rng, train, output_attentions, output_hidden_states, return_dict) and returns: outputs, outputs