class_docstr: A Flax multi-head attention module as described in: https://arxiv.org/abs/1706.03762
  
  Parameters:
      query_dim (:obj:`int`):
          Input hidden states dimension
      heads (:obj:`int`, *optional*, defaults to 8):
          Number of heads
      dim_head (:obj:`int`, *optional*, defaults to 64):
          Hidden states dimension inside each head
      dropout (:obj:`float`, *optional*, defaults to 0.0):
          Dropout rate
      use_memory_efficient_attention (`bool`, *optional*, defaults to `False`):
          enable memory efficient attention https://arxiv.org/abs/2112.05682
      split_head_dim (`bool`, *optional*, defaults to `False`):
          Whether to split the head dimension into a new axis for the self-attention computation. In most cases,
          enabling this flag should speed up the computation for Stable Diffusion 2.x and Stable Diffusion XL.
      dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):
          Parameters `dtype`
code: def __call__(self, hidden_states, context=None, deterministic=True):
      context = hidden_states if context is None else context
  
      query_proj = self.query(hidden_states)
      key_proj = self.key(context)
      value_proj = self.value(context)
  
      if self.split_head_dim:
          b = hidden_states.shape[0]
          query_states = jnp.reshape(
              query_proj, (b, -1, self.heads, self.dim_head)
          )
          key_states = jnp.reshape(
              key_proj, (b, -1, self.heads, self.dim_head)
          )
          value_states = jnp.reshape(
              value_proj, (b, -1, self.heads, self.dim_head)
          )
      else:
          query_states = self.reshape_heads_to_batch_dim(query_proj)
          key_states = self.reshape_heads_to_batch_dim(key_proj)
          value_states = self.reshape_heads_to_batch_dim(value_proj)
  
      if self.use_memory_efficient_attention:
          query_states = query_states.transpose(1, 0, 2)
          key_states = key_states.transpose(1, 0, 2)
          value_states = value_states.transpose(1, 0, 2)
  
          # this if statement create a chunk size for each layer of the unet
          # the chunk size is equal to the query_length dimension of the deepest layer of the unet
  
          flatten_latent_dim = query_states.shape[-3]
          if flatten_latent_dim % 64 == 0:
              query_chunk_size = int(flatten_latent_dim / 64)
          elif flatten_latent_dim % 16 == 0:
              query_chunk_size = int(flatten_latent_dim / 16)
          elif flatten_latent_dim % 4 == 0:
              query_chunk_size = int(flatten_latent_dim / 4)
          else:
              query_chunk_size = int(flatten_latent_dim)
  
          hidden_states = jax_memory_efficient_attention(
              query_states,
              key_states,
              value_states,
              query_chunk_size=query_chunk_size,
              key_chunk_size=4096 * 4,
          )
  
          hidden_states = hidden_states.transpose(1, 0, 2)
      else:
          # compute attentions
          if self.split_head_dim:
              attention_scores = jnp.einsum(
                  "b t n h, b f n h -> b n f t",
                  key_states,
                  query_states,
              )
          else:
              attention_scores = jnp.einsum(
                  "b i d, b j d->b i j", query_states, key_states
              )
  
          attention_scores = attention_scores * self.scale
          attention_probs = nn.softmax(
              attention_scores, axis=-1 if self.split_head_dim else 2
          )
  
          # attend to values
          if self.split_head_dim:
              hidden_states = jnp.einsum(
                  "b n f t, b t n h -> b f n h",
                  attention_probs,
                  value_states,
              )
              b = hidden_states.shape[0]
              hidden_states = jnp.reshape(
                  hidden_states, (b, -1, self.heads * self.dim_head)
              )
          else:
              hidden_states = jnp.einsum(
                  "b i j, b j d -> b i d", attention_probs, value_states
              )
              hidden_states = self.reshape_batch_dim_to_heads(
                  hidden_states
              )
  
      hidden_states = self.proj_attn(hidden_states)
      return self.dropout_layer(
          hidden_states, deterministic=deterministic
      )
  
imports: import functools
  import math
  import flax.linen
  import jax
  import jax.numpy
label_desc: Use image node_id 4 for calling the FlaxAttention obj's underlying member method code with example usage: obj.__call__(hidden_states, context, deterministic) and returns: self