class_docstr: A Flax transformer block layer with `GLU` (Gated Linear Unit) activation function as described in:
  https://arxiv.org/abs/1706.03762
  
  
  Parameters:
      dim (:obj:`int`):
          Inner hidden states dimension
      n_heads (:obj:`int`):
          Number of heads
      d_head (:obj:`int`):
          Hidden states dimension inside each head
      dropout (:obj:`float`, *optional*, defaults to 0.0):
          Dropout rate
      only_cross_attention (`bool`, defaults to `False`):
          Whether to only apply cross attention.
      dtype (:obj:`jnp.dtype`, *optional*, defaults to jnp.float32):
          Parameters `dtype`
      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.
code: def setup(self):
      # self attention (or cross_attention if only_cross_attention is True)
      self.attn1 = FlaxAttention(
          self.dim,
          self.n_heads,
          self.d_head,
          self.dropout,
          self.use_memory_efficient_attention,
          self.split_head_dim,
          dtype=self.dtype,
      )
      # cross attention
      self.attn2 = FlaxAttention(
          self.dim,
          self.n_heads,
          self.d_head,
          self.dropout,
          self.use_memory_efficient_attention,
          self.split_head_dim,
          dtype=self.dtype,
      )
      self.ff = FlaxFeedForward(
          dim=self.dim, dropout=self.dropout, dtype=self.dtype
      )
      self.norm1 = nn.LayerNorm(epsilon=1e-5, dtype=self.dtype)
      self.norm2 = nn.LayerNorm(epsilon=1e-5, dtype=self.dtype)
      self.norm3 = nn.LayerNorm(epsilon=1e-5, dtype=self.dtype)
      self.dropout_layer = nn.Dropout(rate=self.dropout)
  
imports: import functools
  import math
  import flax.linen
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling the FlaxBasicTransformerBlock obj's underlying member method code with example usage: obj.setup() without return types