Use int prng (since we're only testing uints)

jax-ml · Apr 23, 2020 · a30eaec · a30eaec
1 parent 322a04e
commit a30eaec
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/lax_test.py b/tests/lax_test.py
@@ -152,7 +152,7 @@ def op_record(op, nargs, dtypes, rng_factory, tol=None):
     op_record("bitwise_not", 1, bool_dtypes, jtu.rand_small),
     op_record("bitwise_or", 2, bool_dtypes, jtu.rand_small),
     op_record("bitwise_xor", 2, bool_dtypes, jtu.rand_small),
-    op_record("population_count", 1, uint_dtypes, jtu.rand_small),
+    op_record("population_count", 1, uint_dtypes, partial(jtu.rand_int, 1 << 32)),
 
     op_record("add", 2, default_dtypes + complex_dtypes, jtu.rand_small),
     op_record("sub", 2, default_dtypes + complex_dtypes, jtu.rand_small),