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","."]
docstr: Args:
      batch_size (`int`):
          batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.
      max_length (`int`):
          maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized
          cache.
docstr_tok: ["Args",":","batch_size","(","`","int","`",")",":","batch_size","used","for","fast","auto-regressive","decoding",".","Defines","the","batch","size","of","the","initialized","cache",".","max_length","(","`","int","`",")",":","maximum","possible","length","for","auto-regressive","decoding",".","Defines","the","sequence","length","of","the","initialized","cache","."]
code: def init_cache(self, batch_size, max_length):
      r"""
      Args:
          batch_size (`int`):
              batch_size used for fast auto-regressive decoding. Defines the batch size of the initialized cache.
          max_length (`int`):
              maximum possible length for auto-regressive decoding. Defines the sequence length of the initialized
              cache.
      """
      # init input variables to retrieve cache
      input_ids = jnp.ones((batch_size, max_length))
      attention_mask = jnp.ones_like(input_ids)
      position_ids = jnp.broadcast_to(
          jnp.arange(jnp.atleast_2d(input_ids).shape[-1]),
          input_ids.shape,
      )
  
      init_variables = self.module.init(
          jax.random.PRNGKey(0),
          input_ids,
          attention_mask,
          position_ids,
          return_dict=False,
          init_cache=True,
      )
      return unfreeze(init_variables["cache"])
  
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 3 for calling the FlaxLlamaPreTrainedModel obj's underlying member method code with example usage: obj.init_cache(batch_size, max_length) and returns: unfreeze