class_docstr: Transformer pure decoder stack for language modelling.
  
  Args:
    config: TransformerConfig dataclass containing hyperparameters.
docstr: Applies TransformerLM on the inputs.
  
  Args:
    inputs: target data.
    inputs_positions: input subsequence positions for packed examples.
    inputs_segmentation: input segmentation info for packed examples.
  
  Returns:
    logits array from transformer decoder.
docstr_tok: ["Applies","TransformerLM","on","the","inputs",".","Args",":","inputs",":","target","data",".","inputs_positions",":","input","subsequence","positions","for","packed","examples",".","inputs_segmentation",":","input","segmentation","info","for","packed","examples",".","Returns",":","logits","array","from","transformer","decoder","."]
code: def __call__(
      self,
      inputs,
      *,
      inputs_positions=None,
      inputs_segmentation=None,
      rngs: nnx.Rngs | None = None,
  ):
      """Applies TransformerLM on the inputs.
  
      Args:
        inputs: target data.
        inputs_positions: input subsequence positions for packed examples.
        inputs_segmentation: input segmentation info for packed examples.
  
      Returns:
        logits array from transformer decoder.
      """
      config = self.config
  
      # Make padding attention masks.
      if flaglib.flags.get("decode", False):
          # for fast autoregressive decoding we use no decoder mask
          decoder_mask = None
      else:
          decoder_mask = nnx.combine_masks(
              nnx.make_attention_mask(
                  inputs > 0, inputs > 0, dtype=config.dtype
              ),
              nnx.make_causal_mask(inputs, dtype=config.dtype),
          )
  
      # Add segmentation block-diagonal attention masks if using segmented data.
      if inputs_segmentation is not None:
          decoder_mask = nnx.combine_masks(
              decoder_mask,
              nnx.make_attention_mask(
                  inputs_segmentation,
                  inputs_segmentation,
                  jnp.equal,
                  dtype=config.dtype,
              ),
          )
  
      logits = self.decoder(
          inputs,
          inputs_positions=inputs_positions,
          inputs_segmentation=inputs_segmentation,
          decoder_mask=decoder_mask,
          rngs=rngs,
      )
      return logits.astype(self.config.dtype)
  
imports: from __future__ import annotations
  import dataclasses
  from typing import Any, Optional
  import jax
  import jax.numpy
  import numpy
  from jax import lax
  from flax.experimental import nnx
  from flax.experimental.nnx.nnx import flaglib
  from flax.experimental.nnx.examples.lm1b.configs import default
label_desc: Use image node_id 2 for calling the TransformerLM obj's underlying member method code with example usage: obj.__call__(inputs) and returns: logits