class_docstr: Self-attention block of Transformer.
  
  Attributes:
    num_heads: Number of heads.
class_docstr_tok: ["Self-attention","block","of","Transformer",".","Attributes",":","num_heads",":","Number","of","heads","."]
code: def __call__(self, x: jnp.ndarray, attn_mask=None) -> jnp.ndarray:
      xn = LayerNorm(name="ln_1")(x)
      x = x + nn.SelfAttention(
          self.num_heads, name="attn", deterministic=True
      )(xn, attn_mask)
      xn = LayerNorm(name="ln_2")(x)
      x = x + MLP(name="mlp")(xn)
      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 1 for calling the ResidualAttentionBlock obj's underlying member method code with example usage: obj.__call__(x, attn_mask) and returns: x