def run_inner(idx, loop_val):
        """Runs one slice of attention (multiple heads, but one example)."""
        o_all, s_all, i_ct_all, w_ct_all = loop_val
        idx = idx * self._n_parallel_heads
        example_idx = idx // self._n_heads
        head_idx_lo = idx % self._n_heads
        head_range = head_idx_lo + np.arange(n_parallel_heads, dtype=np.int32)
        state_range = idx + np.arange(n_parallel_heads, dtype=np.int32)

        i_mh = fastmath.nested_map(lambda x: x[example_idx], inputs)
        w_mh = fastmath.nested_map(lambda w: w[head_range], weights)
        s_mh = fastmath.nested_map(lambda s: s[state_range], state)
        def forward_unbatched_h(i_h, w_h, s_h):
          return forward_unbatched(*i_h, weights=w_h, state=s_h)
        def forward_fn(i_mh, w_mh):
          o_mh, new_s_mh = fastmath.vmap(
              forward_unbatched_h, in_axes=(None, 0, 0), out_axes=0)(
                  i_mh, w_mh, s_mh)
          o_mh = np.sum(o_mh, axis=0)
          return o_mh, new_s_mh

        if compute_grad:
          o_mh, backward_fn, s_mh = vjp(forward_fn, i_mh, w_mh, has_aux=True)
          ct_mh = output_grad[example_idx]
          assert o_mh.shape == ct_mh.shape
          i_ct_mh, w_ct_mh = backward_fn(ct_mh)
        else:
          o_mh, s_mh = forward_fn(i_mh, w_mh)

        if compute_output:
          o_all = fastmath.index_add(o_all, example_idx, o_mh)
        if update_state:
          s_all = tree_update(s_all, state_range, s_mh)
        if compute_grad:
          i_ct_all = tree_add(i_ct_all, example_idx, i_ct_mh)
          w_ct_all = tree_add(w_ct_all, head_range, w_ct_mh)
        return (o_all, s_all, i_ct_all, w_ct_all)
Runs one slice of attention (multiple heads, but one example).
trax/layers/research/efficient_attention.py