def dense_to_jagged(
    dense: torch.Tensor,
    offsets_list: List[torch.Tensor],
    padding_value: float = 0.0,
) -> torch.Tensor:
    """
    Convert a dense Tensor into a jagged Tensor (using the offsets).

    The function converts a rectangular dense Tensor to a compactly
    represented subset of its values: a jagged Tensor, using the offsets.
    The padding_value is used at the positions of the resulting jagged
    Tensor where the input dense Tensor doesn't have elements.

    Parameters
    ----------
    dense : torch.Tensor
        A Tensor with the shape `[batch_dim, N1, ..., Nn, D1, ..., Dm]`.
        The first n+1 dimensions of the dense Tensor are encoded into
        the first `total_length` dimension of the resulting jagged
        Tensor, using the specified offsets. Importantly, the values in
        the dense Tensor outside of what the offsets specify are omitted
        in the resulting jagged Tensor.

    offsets_list : List[torch.Tensor]
        A list of rank-1 Tensors, each representing the offsets along one
        of the jagged dimensions of the jagged Tensor. The number of offsets
        Tensors in the list must correspond to the number of jagged dimensions
        encoded in the first `total_length` dimension of the resulting jagged
        Tensor. The offsets Tensors must be consistent with the offset
        specification:

            - batch_dim == len(offsets[0]) - 1
            - offsets[i][-1] == len(offsets[i+1])) - 1
            - offsets[-1][-1] == total_length

    padding_value : float
        The value to fill the jagged Tensor with at the positions where
        there are no elements in the dense Tensor (e.g., the consecutive
        offset difference is longer than the corresponding N dimension
        in the dense Tensor input). Default: 0.0.

    Returns
    -------
    torch.Tensor
        The jagged tensor converted from the `dense` Tensor using
        the offsets, with the `padding_value` at the positions
        not available in the `dense` Tensor.
    """
    assert all(t.dim() == 1 for t in offsets_list)
    offsets_list = [list(t.cpu().numpy()) for t in offsets_list]

    _check_offsets(offsets_list)
    assert len(offsets_list) < len(dense.shape)

    total_length = offsets_list[-1][-1]
    inner_shape = dense.shape[1 + len(offsets_list) :]
    jagged_shape = [total_length, *inner_shape]

    dtype = torch_dtype_to_string(dense.dtype)
    result = get_torch_full_tensor(
        shape=jagged_shape,
        fill_value=padding_value,
        dtype=dtype,
    )

    for dense_idx in product(*[range(d) for d in dense.shape[: 1 + len(offsets_list)]]):
        jagged_idx = _dense_idx_to_jagged_idx(
            dense_idx=dense_idx,
            offsets_list=offsets_list,
        )
        if jagged_idx != -1:
            result[jagged_idx] = dense[tuple(dense_idx)]

    return result
Convert a dense Tensor into a jagged Tensor (using the offsets).

The function converts a rectangular dense Tensor to a compactly
represented subset of its values: a jagged Tensor, using the offsets.
The padding_value is used at the positions of the resulting jagged
Tensor where the input dense Tensor doesn't have elements.

Parameters
----------
dense : torch.Tensor
A Tensor with the shape `[batch_dim, N1, ..., Nn, D1, ..., Dm]`.
The first n+1 dimensions of the dense Tensor are encoded into
the first `total_length` dimension of the resulting jagged
Tensor, using the specified offsets. Importantly, the values in
the dense Tensor outside of what the offsets specify are omitted
in the resulting jagged Tensor.

offsets_list : List[torch.Tensor]
A list of rank-1 Tensors, each representing the offsets along one
of the jagged dimensions of the jagged Tensor. The number of offsets
Tensors in the list must correspond to the number of jagged dimensions
encoded in the first `total_length` dimension of the resulting jagged
Tensor. The offsets Tensors must be consistent with the offset
specification:

- batch_dim == len(offsets[0]) - 1
- offsets[i][-1] == len(offsets[i+1])) - 1
- offsets[-1][-1] == total_length

padding_value : float
The value to fill the jagged Tensor with at the positions where
there are no elements in the dense Tensor (e.g., the consecutive
offset difference is longer than the corresponding N dimension
in the dense Tensor input). Default: 0.0.

Returns
-------
torch.Tensor
The jagged tensor converted from the `dense` Tensor using
the offsets, with the `padding_value` at the positions
not available in the `dense` Tensor.
python/aitemplate/testing/jagged_utils.py