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 __init__(
      self,
      config: LlamaConfig,
      input_shape: Tuple = (1, 1),
      seed: int = 0,
      dtype: jnp.dtype = jnp.float32,
      _do_init: bool = True,
      **kwargs,
  ):
      module = self.module_class(config=config, dtype=dtype, **kwargs)
      super().__init__(
          config,
          module,
          input_shape=input_shape,
          seed=seed,
          dtype=dtype,
          _do_init=_do_init,
      )
  
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 1 to create a new FlaxLlamaPreTrainedModel object from inherited base classes: FlaxPreTrainedModel with example: obj = FlaxLlamaPreTrainedModel(config, input_shape, seed, dtype, _do_init)