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