def forward_unbatched(self, x, mask=None, *,
                        weights, state, rng, update_state):
    """Perform attention for a single batch element and head.

    Args:
      x: Inputs for a single example (subclasses may use different inputs)
      mask: Mask for the inputs.
      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._bias:
      if self._share_qk:
        w_q, w_v, w_o, b_q, b_v = weights
      else:
        w_q, w_k, w_v, w_o, b_q, b_k, b_v = weights
    else:
      if self._share_qk:
        w_q, w_v, w_o = weights
      else:
        w_q, w_k, w_v, w_o = weights

    q = np.matmul(x, w_q)
    k = None
    if not self._share_qk:
      k = np.matmul(x, w_k)
    v = np.matmul(x, w_v)

    if self._bias:
      q = q + b_q
      if not self._share_qk:
        k = k + b_k
      v = v + b_v

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

    assert (mask is not None) == self._masked
    if self._masked:
      # mask is a boolean array (True means "is valid token")
      ones_like_mask = np.ones_like(mask, dtype=np.int32)
      kv_info = kv_info * np.where(mask, ones_like_mask, -ones_like_mask)

    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,
        )

    out = np.matmul(o, w_o)
    out = apply_broadcasted_dropout(out, self._output_dropout, output_rng)
    return out, state
Perform attention for a single batch element and head.

    Args:
      x: Inputs for a single example (subclasses may use different inputs)
      mask: Mask for the inputs.
      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.
    
trax/layers/research/efficient_attention.py