class_docstr: An autoregressive transformer-based language model.
class_docstr_tok: ["An","autoregressive","transformer-based","language","model","."]
docstr: Forward pass, producing a sequence of logits.
docstr_tok: ["Forward","pass",",","producing","a","sequence","of","logits","."]
code: def __call__(
      self,
      tokens: jax.Array,  # Batch of sequences of input tokens, shape [B, T].
  ) -> (
      jax.Array
  ):  # Batch of sequences of output token logits, shape [B, T, V].
      """Forward pass, producing a sequence of logits."""
      input_mask = jnp.greater(tokens, self.pad_token)
      unused_batch_size, seq_len = tokens.shape
  
      # Embed the input tokens and positions.
      embed_init = hk.initializers.TruncatedNormal(stddev=0.02)
      token_embedding_map = hk.Embed(
          self.vocab_size, embed_dim=self.model_size, w_init=embed_init
      )
      token_embeddings = token_embedding_map(tokens)
      positional_embeddings = hk.get_parameter(
          "positional_embeddings",
          [seq_len, self.model_size],
          init=embed_init,
      )
      input_embeddings = (
          token_embeddings + positional_embeddings
      )  # [B, T, D]
  
      # Run the transformer over the inputs.
      embeddings = self.transformer(
          input_embeddings, input_mask
      )  # [B, T, D]
  
      # Decode the embeddings (here, we use untied weights).
      return hk.Linear(self.vocab_size)(embeddings)
  
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 LanguageModel obj's underlying member method code with example usage: obj.__call__(tokens) and returns: hk