diff --git a/tests/python/unittest/test_random.py b/tests/python/unittest/test_random.py index 43e9608934e3..2d1bf1284b63 100644 --- a/tests/python/unittest/test_random.py +++ b/tests/python/unittest/test_random.py @@ -447,20 +447,20 @@ def test_uniform_generator(): verify_generator(generator=generator_mx_same_seed, buckets=buckets, probs=probs) @with_seed() -@unittest.skip('Flaky test, tracked in: https://github.com/apache/incubator-mxnet/issues/9856') def test_gamma_generator(): + success_rate = 0.05 ctx = mx.context.current_context() for dtype in ['float16', 'float32', 'float64']: for kappa, theta in [(0.5, 1.0), (1.0, 5.0)]: print("ctx=%s, dtype=%s, Shape=%g, Scale=%g:" % (ctx, dtype, kappa, theta)) buckets, probs = gen_buckets_probs_with_ppf(lambda x: ss.gamma.ppf(x, a=kappa, loc=0, scale=theta), 5) generator_mx = lambda x: mx.nd.random.gamma(kappa, theta, shape=x, ctx=ctx, dtype=dtype).asnumpy() - verify_generator(generator=generator_mx, buckets=buckets, probs=probs) + verify_generator(generator=generator_mx, buckets=buckets, probs=probs, success_rate=success_rate) generator_mx_same_seed = \ lambda x: np.concatenate( [mx.nd.random.gamma(kappa, theta, shape=x // 10, ctx=ctx, dtype=dtype).asnumpy() for _ in range(10)]) - verify_generator(generator=generator_mx_same_seed, buckets=buckets, probs=probs) + verify_generator(generator=generator_mx_same_seed, buckets=buckets, probs=probs, success_rate=success_rate) @with_seed() def test_exponential_generator():