class_docstr: A transformer stack.
class_docstr_tok: ["A","transformer","stack","."]
docstr: Transforms input embedding sequences to output embedding sequences.
docstr_tok: ["Transforms","input","embedding","sequences","to","output","embedding","sequences","."]
code: def __call__(
      self,
      embeddings: jax.Array,  # [B, T, D]
      mask: jax.Array,  # [B, T]
  ) -> jax.Array:  # [B, T, D]
      """Transforms input embedding sequences to output embedding sequences."""
  
      initializer = hk.initializers.VarianceScaling(2 / self.num_layers)
      _, seq_len, model_size = embeddings.shape
  
      # Compute causal mask for autoregressive sequence modelling.
      mask = mask[:, None, None, :]  # [B, H=1, T'=1, T]
      causal_mask = np.tril(
          np.ones((1, 1, seq_len, seq_len))
      )  # [B=1, H=1, T, T]
      mask = mask * causal_mask  # [B, H=1, T, T]
  
      h = embeddings
      for _ in range(self.num_layers):
          # First the attention block.
          attn_block = hk.MultiHeadAttention(
              num_heads=self.num_heads,
              key_size=self.attn_size,
              model_size=model_size,
              w_init=initializer,
          )
          h_norm = _layer_norm(h)
          h_attn = attn_block(h_norm, h_norm, h_norm, mask=mask)
          h_attn = hk.dropout(
              hk.next_rng_key(), self.dropout_rate, h_attn
          )
          h = h + h_attn
  
          # Then the dense block.
          dense_block = hk.Sequential(
              [
                  hk.Linear(
                      self.widening_factor * model_size,
                      w_init=initializer,
                  ),
                  jax.nn.gelu,
                  hk.Linear(model_size, w_init=initializer),
              ]
          )
          h_norm = _layer_norm(h)
          h_dense = dense_block(h_norm)
          h_dense = hk.dropout(
              hk.next_rng_key(), self.dropout_rate, h_dense
          )
          h = h + h_dense
  
      return _layer_norm(h)
  
file_path: dm-haiku/examples/transformer/model.py
imports: import dataclasses
  from typing import Optional
  import haiku
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling the Transformer obj's underlying member method code with example usage: obj.__call__(embeddings, mask) and returns: _layer_norm