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

        i_h = fastmath.nested_map(lambda x: x[example_idx], inputs)
        w_h = fastmath.nested_map(lambda w: w[head_idx], weights)
        s_h = fastmath.nested_map(lambda s: s[idx], state)

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

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

        if compute_output:
          o_all = fastmath.index_add(o_all, example_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, example_idx, i_ct_h)
          w_ct_all = tree_add(w_ct_all, head_idx, w_ct_h)
        return (o_all, s_all, i_ct_all, w_ct_all)
Runs one slice of attention (for a single head).
trax/layers/research/efficient_attention.py