def _incremental_forward_unbatched(self, *inputs, q_start, q_len,
                                     weights, state):
    """Perform fast inference for a single batch element and head.

    Subclasses should override this method.

    Args:
      *inputs: Inputs for a single example (subclasses may use different inputs)
      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.
    Returns:
      A tuple (output, new_state) -- output and new state for a single example
      and attention head.
    """
    raise NotImplementedError(
        'Fast inference is not implemented for this attention type.')
Perform fast inference for a single batch element and head.

    Subclasses should override this method.

    Args:
      *inputs: Inputs for a single example (subclasses may use different inputs)
      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.
    Returns:
      A tuple (output, new_state) -- output and new state for a single example
      and attention head.
    
trax/layers/research/efficient_attention.py