class_docstr: Transformer pure decoder stack for language modelling.
  
  Args:
    config: TransformerConfig dataclass containing hyperparameters.
class_docstr_tok: ["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
  ):
      """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 config.decode:
          # for fast autoregressive decoding we use no decoder mask
          decoder_mask = None
      else:
          decoder_mask = nn.combine_masks(
              nn.make_attention_mask(
                  inputs > 0, inputs > 0, dtype=config.dtype
              ),
              nn.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 = nn.combine_masks(
              decoder_mask,
              nn.make_attention_mask(
                  inputs_segmentation,
                  inputs_segmentation,
                  jnp.equal,
                  dtype=config.dtype,
              ),
          )
  
      logits = Decoder(
          config=config, shared_embedding=None, name="decoder"
      )(
          inputs,
          inputs_positions=inputs_positions,
          inputs_segmentation=inputs_segmentation,
          decoder_mask=decoder_mask,
          encoder_decoder_mask=None,
      )
      return logits.astype(self.config.dtype)
  
imports: from __future__ import annotations
  from typing import Callable, Any, Optional
  from flax import linen
  from flax import struct
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling the TransformerLM obj's underlying member method code with example usage: obj.__call__(inputs, inputs_positions, inputs_segmentation) and returns: logits