class_docstr: Transformer module.
  
  Attributes:
    features: Number of features.
    num_layers: Number of layers for each block.
    num_heads: Number of heads.
    use_underscore_module_name: Optionally replace '.' with '_' in parameter
      naming for PAX checkpoint loading.
class_docstr_tok: ["Transformer","module",".","Attributes",":","features",":","Number","of","features",".","num_layers",":","Number","of","layers","for","each","block",".","num_heads",":","Number","of","heads",".","use_underscore_module_name",":","Optionally","replace","'",".","'","with","'","_","'","in","parameter","naming","for","PAX","checkpoint","loading","."]
code: def __call__(
      self, x: jnp.ndarray, attn_mask: Optional[jnp.ndarray] = None
  ) -> jnp.ndarray:
      def _n(name):
          """A helper function that optionally replace '.' with '_'."""
          if self.use_underscore_module_name:
              return name.replace(".", "_")
          else:
              return name
  
      for i in range(self.num_layers):
          x = ResidualAttentionBlock(
              num_heads=self.num_heads, name=_n(f"resblocks.{i}")
          )(x, attn_mask)
      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 Transformer obj's underlying member method code with example usage: obj.__call__(x, attn_mask) and returns: x, name, name