code: def setup(self):
      config = self.config
      self.embed_dim = config.hidden_size
      self.num_heads = config.num_attention_heads
      self.head_dim = self.embed_dim // self.num_heads
      self.attention_softmax_in_fp32 = self.dtype is not jnp.float32
  
      dense = partial(
          nn.Dense,
          self.embed_dim,
          use_bias=config.attention_bias,
          dtype=self.dtype,
          kernel_init=jax.nn.initializers.normal(
              self.config.initializer_range
          ),
      )
  
      self.q_proj, self.k_proj, self.v_proj = dense(), dense(), dense()
      self.o_proj = dense()
  
      self.causal_mask = make_causal_mask(
          jnp.ones((1, config.max_position_embeddings), dtype="bool"),
          dtype="bool",
      )
      self.rotary_emb = FlaxLlamaRotaryEmbedding(
          config, dtype=self.dtype
      )
  
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 for calling the FlaxLlamaAttention obj's underlying member method code with example usage: obj.setup() without return types