def generate_offsets(
    batch_size: int,
    max_seq_len: int,
    load_factor: float,
    offsets_dtype: str,
    spread_radius: float = 0.1,
) -> torch.Tensor:
    """
    Generate a rank-1 Tensor of offsets for the given load factor.

    This function generates a single linear offset Tensor for a
    single jagged dimension in a jagged Tensor with the batch_dim
    equal to `batch_size` and maximum value along the jagged
    dimension equal to `max_seq_len`. The `load_factor` in [0, 1]
    specifies how "full" should the jagged Tensor described by
    the resulting offsets should be, compared to the corresponding
    dense Tensor with a rectangular shape [batch_size, max_seq_len,
    D1, ..., Dm]. The offset differences (== the lengths along the
    jagged dimensions) are sampled randomly, to arrive close (but not
    necessarily equal) to the specified `load_factor` in total.

    When sampled out of the [0, N] interval, the offset differences
    are clamed to stay within the [0, N] interval.

    Parameters
    ----------
    batch_size : int
        The batch_dim of the jagged Tensor specified by the offsets.
    max_seq_len : int
        The maximum length along the jagged dimension specified by
        the offsets.
    load_factor : float
        The fraction of the [batch_size, max_seq_len, D1, ..., Dm]-
        shaped dense Tensor that the total (compactly represented)
        jagged Tensor data should correspond to.
    offsets_dtype : str
        The type of the resulting offsets Tensor.
    spread_radius : float
        The radius of the spread around int(max_seq_len * load_factor)
        that the offset differences should be randomly sampled from.
        Default: 0.1.

    Returns
    -------
    torch.Tensor
        The resulting rank-1 Tensor of offsets.
    """
    assert 0 <= load_factor <= 1
    assert 0 <= spread_radius <= 1

    if load_factor < 1:
        spread = int(max_seq_len * spread_radius)
        mean = int(max_seq_len * load_factor)
        lengths = [
            mean + random.randint(-spread, spread + 1) for _ in range(batch_size)
        ]
        lengths = [max(min(L, max_seq_len), 0) for L in lengths]
    else:
        lengths = [max_seq_len] * batch_size

    offsets = [0]
    for length in lengths:
        offsets.append(offsets[-1] + length)

    torch_offsets_dtype = string_to_torch_dtype(offsets_dtype)
    return torch.tensor(offsets, dtype=torch_offsets_dtype).cuda()
Generate a rank-1 Tensor of offsets for the given load factor.

This function generates a single linear offset Tensor for a
single jagged dimension in a jagged Tensor with the batch_dim
equal to `batch_size` and maximum value along the jagged
dimension equal to `max_seq_len`. The `load_factor` in [0, 1]
specifies how "full" should the jagged Tensor described by
the resulting offsets should be, compared to the corresponding
dense Tensor with a rectangular shape [batch_size, max_seq_len,
D1, ..., Dm]. The offset differences (== the lengths along the
jagged dimensions) are sampled randomly, to arrive close (but not
necessarily equal) to the specified `load_factor` in total.

When sampled out of the [0, N] interval, the offset differences
are clamed to stay within the [0, N] interval.

Parameters
----------
batch_size : int
The batch_dim of the jagged Tensor specified by the offsets.
max_seq_len : int
The maximum length along the jagged dimension specified by
the offsets.
load_factor : float
The fraction of the [batch_size, max_seq_len, D1, ..., Dm]-
shaped dense Tensor that the total (compactly represented)
jagged Tensor data should correspond to.
offsets_dtype : str
The type of the resulting offsets Tensor.
spread_radius : float
The radius of the spread around int(max_seq_len * load_factor)
that the offset differences should be randomly sampled from.
Default: 0.1.

Returns
-------
torch.Tensor
The resulting rank-1 Tensor of offsets.
python/aitemplate/testing/jagged_utils.py