class_docstr: Self-attention block of Transformer.
  
  Attributes:
    num_heads: Number of heads.
    droplayer_p: Layer drop probability.
class_docstr_tok: ["Self-attention","block","of","Transformer",".","Attributes",":","num_heads",":","Number","of","heads",".","droplayer_p",":","Layer","drop","probability","."]
code: def __call__(
      self,
      x: jnp.ndarray,
      attn_mask: Optional[jnp.ndarray] = None,
      *,
      deterministic: bool = True,
  ) -> jnp.ndarray:
      xn = LayerNorm(name="ln_0")(x)
      y = nn.SelfAttention(
          self.num_heads, name="attn", deterministic=deterministic
      )(xn, attn_mask)
  
      # Droplayer.
      drop_pattern = self.get_drop_pattern(y, deterministic)
      x = y * (1.0 - drop_pattern) + x
  
      xn = LayerNorm(name="ln_1")(x)
      y = MLP(name="mlp")(xn)
  
      # Droplayer.
      drop_pattern = self.get_drop_pattern(x, deterministic)
      x = y * (1.0 - drop_pattern) + x
      return x
  
imports: import functools
  from typing import Sequence, Optional, Union, Tuple
  import flax.linen
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling the ResidualAttentionBlock obj's underlying member method code with example usage: obj.__call__(x, attn_mask) and returns: x