def run_inner(idx, loop_val):
        """Runs one slice of attention (all heads for one or more examples)."""
        o_all, s_all, i_ct_all, w_ct_all = loop_val
        idx = idx * n_parallel_heads
        example_idx_lo = idx // self._n_heads
        example_range = example_idx_lo + np.arange(
            n_parallel_heads // self._n_heads, dtype=np.int32)
        state_range = idx + np.arange(n_parallel_heads, dtype=np.int32)

        i_mex = fastmath.nested_map(lambda x: x[example_range], inputs)
        s_mex = fastmath.nested_map(
            lambda s: np.reshape(s[state_range],  # pylint: disable=g-long-lambda
                                 (-1, self._n_heads) + s.shape[1:]),
            state)
        def forward_fn(i_mex, w_all):
          o_mex, new_s_mex = fastmath.vmap(
              forward_single_example, in_axes=(0, None, 0), out_axes=(0, 0))(
                  i_mex, w_all, s_mex)
          new_s_mex = fastmath.nested_map(
              lambda s: np.reshape(s, (n_parallel_heads,) + s.shape[2:]),
              new_s_mex)
          return o_mex.astype(i_mex[0].dtype), new_s_mex

        if compute_grad:
          o_mex, backward_fn, s_mex = vjp(forward_fn, i_mex, weights,
                                          has_aux=True)
          ct_mex = output_grad[example_range]
          assert o_mex.shape == ct_mex.shape, str(ct_mex.shape)
          assert o_mex.dtype == ct_mex.dtype, str(ct_mex.dtype)
          i_ct_mex, w_ct_mex = backward_fn(ct_mex)
        else:
          o_mex, s_mex = forward_fn(i_mex, weights)

        if compute_output:
          o_all = fastmath.index_add(o_all, jax.numpy.index_exp[example_range],
                                     o_mex)
        if update_state:
          s_all = tree_update(s_all, state_range, s_mex)
        if compute_grad:
          i_ct_all = tree_update(i_ct_all, example_range, i_ct_mex)
          w_ct_all = fastmath.nested_map_multiarg(
              lambda old_all, delta_all: old_all + delta_all,
              w_ct_all, w_ct_mex)
        return (o_all, s_all, i_ct_all, w_ct_all)
Runs one slice of attention (all heads for one or more examples).
trax/layers/research/efficient_attention.py