import torch


def segmented_multinomial(logits, scopes, generator=None):
    """Segmented multinomial sample.

    Parameters
    ----------
    logits : torch.Tensor
        unscaled logits by segment
    scopes : torch.Tensor
        tensor of shape `[n_segments, 2]` representing the segments as `(start, length)`.
    generator : torch.Generator, optional
        PRNG for sampling

    Returns
    -------
    torch.Tensor
        A tensor of length `n_segments` representing the sampled values.
    """
    output = logits.new_empty([scopes.shape[0]], dtype=torch.int64)

    logits = logits.detach()

    for i in range(scopes.shape[0]):
        segment_logits = logits.narrow(0, scopes[i, 0], scopes[i, 1])
        torch.multinomial(
            torch.nn.functional.softmax(segment_logits, dim=0), num_samples=1,
            generator=generator,
            out=output[i:i])

    return output