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

    The function converts a jagged Tensor (and the offsets) to
    a rectangular dense Tensor, using the padding_value at the
    positions of the resulting dense Tensor where the input
    jagged Tensor doesn't have elements.

    Parameters
    ----------
    jagged : torch.Tensor
        The jagged Tensor with the shape `[total_length, D1, ..., Dm]`,
        The first dimension, `total_length`, encodes the `batch_dim` and
        the jagged dims of the jagged Tensor. The following m dimensions,
        `D1, ..., Dm`, are regular dense dimensions following the jagged
        dimensions of the 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 `jagged`. 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

    dense_shape : List[int]
        The shape of the resulting dense Tensor. The last m dimensions in
        the `dense_shape` must be equal to `[D1, ..., Dm]` in the jagged
        Tensor shape. The first dimension must be the `batch_dim`. The
        following n dimensions must correspond to the n jagged dimensions
        of the jagged Tensor, with the values equal to the maximum possible
        values of the jagged dimensions.

    padding_value : float
        The value to fill the dense Tensor with at the positions where
        there are no elements in the jagged Tensor. Default: 0.0.

    Returns
    -------
    torch.Tensor
        The dense tensor with the `dense_shape` converted from the
        `jagged` Tensor, with the `padding_value` at other positions.
    """
    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(dense_shape) - len(jagged.shape) == len(offsets_list)
    assert jagged.shape[1:] == tuple(dense_shape[1 + len(offsets_list) :])
    for i, offsets in enumerate(offsets_list):
        dense_dim = dense_shape[i + 1]
        for j in range(1, len(offsets)):
            assert offsets[j] - offsets[j - 1] <= dense_dim

    dtype = torch_dtype_to_string(jagged.dtype)
    result = get_torch_full_tensor(
        shape=dense_shape,
        fill_value=padding_value,
        dtype=dtype,
    )

    total_length = jagged.shape[0]
    for jagged_idx in range(total_length):
        dense_idx = _jagged_idx_to_dense_idx(
            jagged_idx=jagged_idx,
            offsets_list=offsets_list,
        )
        result[tuple(dense_idx)] = jagged[jagged_idx]

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

The function converts a jagged Tensor (and the offsets) to
a rectangular dense Tensor, using the padding_value at the
positions of the resulting dense Tensor where the input
jagged Tensor doesn't have elements.

Parameters
----------
jagged : torch.Tensor
The jagged Tensor with the shape `[total_length, D1, ..., Dm]`,
The first dimension, `total_length`, encodes the `batch_dim` and
the jagged dims of the jagged Tensor. The following m dimensions,
`D1, ..., Dm`, are regular dense dimensions following the jagged
dimensions of the 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 `jagged`. 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

dense_shape : List[int]
The shape of the resulting dense Tensor. The last m dimensions in
the `dense_shape` must be equal to `[D1, ..., Dm]` in the jagged
Tensor shape. The first dimension must be the `batch_dim`. The
following n dimensions must correspond to the n jagged dimensions
of the jagged Tensor, with the values equal to the maximum possible
values of the jagged dimensions.

padding_value : float
The value to fill the dense Tensor with at the positions where
there are no elements in the jagged Tensor. Default: 0.0.

Returns
-------
torch.Tensor
The dense tensor with the `dense_shape` converted from the
`jagged` Tensor, with the `padding_value` at other positions.
python/aitemplate/testing/jagged_utils.py