Function _multinomial

imperative/python/megengine/random/rng.py:213–246

(
    input: Tensor, num_samples: int, replacement: bool, seed: int, handle: int
)

def _multinomial(
    input: Tensor, num_samples: int, replacement: bool, seed: int, handle: int
) -> Tensor:
    handle_cn = None if handle == 0 else _get_rng_handle_compnode(handle)
    assert num_samples > 0, "Multinomial is not defined when num_samples <= 0"
    assert (
        input.ndim == 1 or input.ndim == 2
    ), "Multinomial is not defined when ndim of input is not 1 or 2"
    assert input.shape[-1] != 0, "Multinomial is not defined when num of input == 0"
    assert (
        input.dtype.name == "float32" or input.dtype.name == "float16"
    ), "Multinomial is not defined when input dtype is not float"
    assert (
        replacement or num_samples <= input.shape[-1]
    ), "Multinomial is not defined when num_samples > input.shape[-1] in case of no replacement"

    require_one_dim = False
    if input.ndim == 1:
        input = input.reshape((1, input.size))
        require_one_dim = True
    if replacement:
        input = input / input.sum(-1, True)
    assert (
        handle_cn is None or handle_cn == input.device
    ), "The input ({}) must be the same device with handle ({})".format(
        input.device, handle_cn
    )
    op = MultinomialRNG(
        seed=seed, num_samples=num_samples, replacement=replacement, handle=handle
    )
    (output,) = apply(op, input)
    if require_one_dim:
        output = output.reshape((output.size,))
    return output


def _permutation(n: int, seed: int, device: str, handle: int, dtype: str) -> Tensor: