Method __call__

transformer_layer.py:119–485

Computes attention over a sequence of inputs. Args: xs: input sequence of shape (batch_size, sequence_length, num_hidden) start_of_sequence: An input array of shape (batch_size) --- The following must be passed by keyword only. --- importance: Array of shape (batch_si

(
      self,
      xs: Array,
      start_of_sequence: Array,
      *,
      importance: Optional[Array] = None,
      cross_attention_kv: Optional[Tuple[Array, Array]] = None,
      window_state: Optional[WindowState] = None,
      decoder_state: Optional[DecoderState] = None,
  )

    return (DecoderState(out_decoder_state), out_keys, out_values)

  def __call__(
      self,
      xs: Array,
      start_of_sequence: Array,
      *,
      importance: Optional[Array] = None,
      cross_attention_kv: Optional[Tuple[Array, Array]] = None,
      window_state: Optional[WindowState] = None,
      decoder_state: Optional[DecoderState] = None,
  ):
    """Computes attention over a sequence of inputs.

    Args:
      xs: input sequence of shape (batch_size, sequence_length, num_hidden)
      start_of_sequence: An input array of shape (batch_size)  --- The following
        must be passed by keyword only. ---
      importance: Array of shape (batch_size, sequence_length). An importance
        bias for attention.
      cross_attention_kv: Keys and values from encoder for cross-attention.
      window_state: State object which contains context from the prior window
        when using a transformer-XL or sliding window. Initially created with
        load_window_state().
      decoder_state: State object for autoregressive decoding, initially created
        with from init_decoder_state().

    Returns:
      (ys: outputs of shape (batch_size, sequence_length, num_hidden),
       importance_score: importance score for the next layer,
       next_window_state: state to pass to the next window,
       next_decoder_state: next decoder state for autoregressive decoding,
       viz_dict: dictionary of visualizations
      )
    """

    xs = jnp.asarray(xs, dtype=self.dtype)
    logging.info("tlayer: recurrent = %r", self.recurrent_attention)
    logging.info("tlayer: compute_importance = %r", self.compute_importance)

    is_training = self.mode == "train"

    # Compute keys, values and queries.
    # ---------------------------------
    logging.info("tlayer: compute keys,values,queries.")
    (keys, values, queries, queries2) = self.tbase.kvq(xs)
    attention_scale_factors = self.tbase.attention_scale_factors()
    (_, sequence_length, num_heads, _) = queries.shape  # (b, k, h, d)

    # Get biases and masks that are shared across windows.
    # ----------------------------------------------------
    if decoder_state is not None:
      logging.info("tlayer: using autoregressive decoder.")
      # When decoding, prior keys,values are loaded from the decoder state.
      # Other values are precomputed, and loaded from the decoder state.
      # The decoder state will be updated with the current token.
      assert window_state is None

      prev_kvi = None
      recurrent_state = None  # Use precomputed recurrent_kvq.