def _create_global_aggregates(hidden_states: np.ndarray, block_ids: np.ndarray, global_seq_len: int) -> np.ndarray:
    """Compute individual block aggregates by summing over individual blocks."""
    # (batch..., seq_len, global_seq_len))
    one_hot_block_ids = jax.nn.one_hot(block_ids, global_seq_len)
    return jnp.einsum("...nd,...ng->...gd", hidden_states, one_hot_block_ids)
Compute individual block aggregates by summing over individual blocks.
dpkd/transformers/src/transformers/models/longt5/modeling_flax_longt5.py