Function test_rng_empty_tensor

imperative/python/test/unit/random/test_rng.py:798–871

(is_symbolic)

@pytest.mark.parametrize("is_symbolic", [None, False, True])
def test_rng_empty_tensor(is_symbolic):
    set_global_seed(1024)
    shapes = [
        (0,),
        (0, 0, 0),
        (10, 0, 10),
    ]

    def fn(shape):
        o1 = random.uniform(0, 1, shape)
        o2 = random.normal(0, 1, shape)
        o3 = random.gamma(2, 1, shape)
        o4 = random.beta(2, 1, shape)
        o5 = random.poisson(2, shape)
        o6 = random.exponential(1.5, shape)
        return o1, o2, o3, o4, o5, o6

    for shape in shapes:
        if is_symbolic is not None:
            fn_ = jit.trace(symbolic=is_symbolic)(fn)
        else:
            fn_ = fn
        for _ in range(3):
            outs = fn_(shape)
            for out in outs:
                np.testing.assert_equal(out.numpy().shape, shape)
            if is_symbolic is None:
                break

    def fn2(n):
        return random.permutation(n=n)

    if is_symbolic is not None:
        fn2 = jit.trace(symbolic=is_symbolic)(fn2)

    for _ in range(3):
        out = fn2(0)
        np.testing.assert_equal(out.numpy().shape, (0,))
        if is_symbolic is None:
            break

    input_shapes_for_multinomial = [
        (0,),
        (0, 0),
        (0, 10),
    ]

    output_shapes_for_multinomial = [
        (0,),
        (0, 0),
        (0, 2),
    ]

    def fn3(shape, replacement):
        input = random.normal(0, 1, shape)
        out = random.multinomial(input=input, num_samples=2, replacement=True)
        return out