(self, x, mask=None, *,
                                     q_start, q_len,
                                     weights, state, rng, update_state)
return out, state
Perform fast inference for a single batch element and head.

    Args:
      x: Inputs for a single example (subclasses may use different inputs)
      mask: inputs mask.
      q_start: Index along the sequence-length dimension that points to the
        first input element that should be used as a query (and not just a key).
      q_len: Number of new query elements in this call to the attention
        mechanism. This is typically 1 for autoregressive decoding, but may be
        longer if initializing a language model with a prefix.
      weights: Weights for a single attention head
      state: State for a single example & attention head pair.
      rng: PRNG key for the layer (shared across all examples and heads)
      update_state: bool: whether to return an updated layer state.

    Returns:
      A tuple (output, new_state) -- output and new state for a single example
      and attention head.
Perform fast inference for a single batch element and head.
Perform fast inference for a single batch element and head .
def _incremental_forward_unbatched(self, x, mask=None, *,
                                     q_start, q_len,
                                     weights, state, rng, update_state):
    """Perform fast inference for a single batch element and head.

    Args:
      x: Inputs for a single example (subclasses may use different inputs)
      mask: inputs mask.
      q_start: Index along the sequence-length dimension that points to the
        first input element that should be used as a query (and not just a key).
      q_len: Number of new query elements in this call to the attention
        mechanism. This is typically 1 for autoregressive decoding, but may be
        longer if initializing a language model with a prefix.
      weights: Weights for a single attention head
      state: State for a single example & attention head pair.
      rng: PRNG key for the layer (shared across all examples and heads)
      update_state: bool: whether to return an updated layer state.

    Returns:
      A tuple (output, new_state) -- output and new state for a single example
      and attention head.
    """
    del update_state
    attend_rng, output_rng = fastmath.random.split(rng)
    if self._share_qk:
      w_q, w_v, w_o = weights
    else:
      w_q, w_k, w_v, w_o = weights

    q_range = q_start + np.arange(q_len, dtype=np.int32)
    if q_len == 1:
      # On TPU, np.matmul(a[:1], b) and np.matmul(a, b)[:1] are not
      # floating-point equivalent, at least in non-jitted code. We correct the
      # discrepancy by duplicating the slice. Floating-point noise may not be
      # an issue when using models, but it makes it harder to write tests that
      # compare fast and slow inference code for equivalence.
      q = np.matmul(np.concatenate([x[q_range]] * 2, 0), w_q)
    else:
      q = np.matmul(x[q_range], w_q)
    if self._share_qk:
      k = length_normalized(np.matmul(x, w_q))
    else:
      k = np.matmul(x, w_k)
    v = np.matmul(x, w_v)

    mask_fn = functools.partial(
        mask_self_attention,
        causal=self._causal, exclude_self=self._share_qk, masked=self._masked)
    q_info = q_range
    kv_info = np.arange(k.shape[-2], dtype=np.int32)

    if self._chunk_len is not None and q_len > self._chunk_len:
      assert q_start == 0
      assert q_len % self._chunk_len == 0
      o, _ = attend(
          q, k, v,
          q_chunk_len=self._chunk_len,
          kv_chunk_len=self._chunk_len,
          n_chunks_before=self._n_chunks_before,
          n_chunks_after=self._n_chunks_after,
          mask_fn=mask_fn, q_info=q_info, kv_info=kv_info,
          dropout=self._attention_dropout, rng=attend_rng,
          )
    else:
      o, _ = attend(
          q, k, v,
          mask_fn=mask_fn, q_info=q_info, kv_info=kv_info,
          dropout=self._attention_dropout, rng=attend_rng,
          )

    out = np.matmul(o, w_o)
    if q_len == 1:
      out = out[:1]
    out = apply_broadcasted_dropout(out, self._output_dropout, output_rng)
    return out, state
def _incremental_forward_unbatched ( self , x , mask = None , * , q_start , q_len , weights , state , rng , update_state ) : del update_state attend_rng , output_rng = fastmath . random . split ( rng ) if self . _share_qk : w_q , w_v , w_o = weights else : w_q , w_k , w_v , w_o = weights q_range = q_start + np . arange ( q_len , dtype = np . int32 ) if q_len == 1 : # On TPU, np.matmul(a[:1], b) and np.matmul(a, b)[:1] are not # floating-point equivalent, at least in non-jitted code. We correct the # discrepancy by duplicating the slice. Floating-point noise may not be # an issue when using models, but it makes it harder to write tests that # compare fast and slow inference code for equivalence. q = np . matmul ( np . concatenate ( [ x [ q_range ] ] * 2 , 0 ) , w_q ) else : q = np . matmul ( x [ q_range ] , w_q ) if self . _share_qk : k = length_normalized ( np . matmul ( x , w_q ) ) else : k = np . matmul ( x , w_k ) v = np . matmul ( x , w_v ) mask_fn = functools . partial ( mask_self_attention , causal = self . _causal , exclude_self = self . _share_qk , masked = self . _masked ) q_info = q_range kv_info = np . arange ( k . shape [ - 2 ] , dtype = np . int32 ) if self . _chunk_len is not None and q_len > self . _chunk_len : assert q_start == 0 assert q_len % self . _chunk_len == 0 o , _ = attend ( q , k , v , q_chunk_len = self . _chunk_len , kv_chunk_len = self . _chunk_len , n_chunks_before = self . _n_chunks_before , n_chunks_after = self . _n_chunks_after , mask_fn = mask_fn , q_info = q_info , kv_info = kv_info , dropout = self . _attention_dropout , rng = attend_rng , ) else : o , _ = attend ( q , k , v , mask_fn = mask_fn , q_info = q_info , kv_info = kv_info , dropout = self . _attention_dropout , rng = attend_rng , ) out = np . matmul ( o , w_o ) if q_len == 1 : out = out [ : 1 ] out = apply_broadcasted_dropout ( out , self . _output_dropout , output_rng ) return out , state