def run_inner(idx, loop_val):
      """Runs one slice of attention (for a single head)."""
      o_all, s_all, i_ct_all = loop_val
      example_idx = idx // self._n_heads
      unused_head_idx = idx % self._n_heads

      s_h = fastmath.nested_map(lambda s: s[idx], state)

      if self._masked:
        i_h = (qk[idx], v[idx], mask[example_idx])
      else:
        i_h = (qk[idx], v[idx])

      def forward_fn(i_h):
        return forward_unbatched(
            *i_h, state=fastmath.stop_gradient(s_h))

      if compute_grad:
        o_h, backward_fn, s_h = vjp(forward_fn, i_h, has_aux=True)
        ct_h = output_grad[idx]
        assert o_h.shape == ct_h.shape
        i_ct_h, = backward_fn(ct_h)
      else:
        o_h, s_h = forward_fn(i_h)

      if compute_output:
        o_all = tree_update(o_all, idx, o_h)
      if update_state:
        s_all = tree_update(s_all, idx, s_h)
      if compute_grad:
        i_ct_all = tree_add(i_ct_all, idx, i_ct_h)
      return (o_all, s_all, i_ct_all)
Runs one slice of attention (for a single head).
trax/layers/research/efficient_attention.py