@triton.jit
def sparse_map_kernel(top_experts_ptr, expert_ends_ptr,
    expert_pad_begins_ptr, sparse_rows_ptr, num_sparse_rows: tl.constexpr,
    num_experts: tl.constexpr, pad_to_multiple: tl.constexpr, block_size:
    tl.constexpr, block_size_expert: tl.constexpr, dtype: tl.constexpr):
    """
    Since the methods we want (histogram, argsort) are not readily available in triton,
    we use a one-hot representation to get the quantities we want.
    TODO: Next triton release will support tl.histogram, maybe argsort.
    """
    block_range = tl.arange(0, block_size)
    expert_range = tl.arange(0, block_size_expert)
    expert_mask = (None if block_size_expert == num_experts else 
        expert_range < num_experts)
    if num_sparse_rows >= block_size:
        expert_index = tl.load(top_experts_ptr + block_range)
    else:
        expert_index = tl.load(top_experts_ptr + block_range, mask=
            block_range < num_sparse_rows, other=num_experts)
    expert_counts = tl.sum((expert_index[:, None] == expert_range[None, :])
        .to(dtype), 0)
    for i in range(1, tl.cdiv(num_sparse_rows, block_size)):
        block_range += block_size
        if num_sparse_rows % block_size == 0:
            expert_index = tl.load(top_experts_ptr + block_range)
        else:
            expert_index = tl.load(top_experts_ptr + block_range, mask=
                block_range < num_sparse_rows, other=num_experts)
        expert_counts += tl.sum((expert_index[:, None] == expert_range[None,
            :]).to(dtype), 0)
    if pad_to_multiple is None:
        expert_counts_padded = expert_counts
    else:
        expert_counts_padded = (expert_counts + pad_to_multiple - 1
            ) // pad_to_multiple * pad_to_multiple
    expert_ends = tl.cumsum(expert_counts_padded)
    expert_begins = expert_ends - expert_counts_padded
    if expert_ends_ptr is not None:
        tl.store(expert_ends_ptr + expert_range, expert_ends, mask=expert_mask)
    if expert_pad_begins_ptr is not None:
        tl.store(expert_pad_begins_ptr + expert_range, expert_begins +
            expert_counts, mask=expert_mask)
    if sparse_rows_ptr is not None:
        block_range = tl.arange(0, block_size)
        for i in range(tl.cdiv(num_sparse_rows, block_size)):
            if num_sparse_rows % block_size == 0:
                mask = None
                expert_index = tl.load(top_experts_ptr + block_range)
            else:
                mask = block_range < num_sparse_rows
                expert_index = tl.load(top_experts_ptr + block_range, mask=
                    mask, other=num_experts)
            expert_one_hot = (expert_index[:, None] == expert_range[None, :]
                ).to(dtype)
            expert_offsets = (tl.cumsum(expert_one_hot, 0) + expert_begins[
                None, :]) * expert_one_hot
            tl.store(sparse_rows_ptr + block_range, tl.sum(expert_offsets, 
                1) - 1, mask=mask)
            expert_begins += tl.sum(expert_one_hot, 0)
            block_range += block_size
