Refactor remaining tests to use NumPy Generator

Author: kc611

tests/compile/function/test_pfunc.py

@@ -64,15 +64,15 @@ def test_doc(self):
     def test_shared(self):
 
         # CHECK: two functions (f1 and f2) can share w
-        w = shared(np.random.rand(2, 2), "w")
+        w = shared(np.random.random((2, 2)), "w")
         wval = w.get_value(borrow=False)
 
         x = dmatrix()
         out1 = w + x
         out2 = w * x
         f1 = pfunc([x], [out1])
         f2 = pfunc([x], [out2])
-        xval = np.random.rand(2, 2)
+        xval = np.random.random((2, 2))
         assert np.all(f1(xval) == xval + wval)
         assert np.all(f2(xval) == xval * wval)
 
@@ -89,7 +89,7 @@ def test_shared(self):
 
     def test_no_shared_as_input(self):
         # Test that shared variables cannot be used as function inputs.
-        w_init = np.random.rand(2, 2)
+        w_init = np.random.random((2, 2))
         w = shared(w_init.copy(), "w")
         with pytest.raises(
             TypeError, match=r"^Cannot use a shared variable \(w\) as explicit input"
@@ -100,8 +100,8 @@ def test_default_container(self):
         # Ensure it is possible to (implicitly) use a shared variable in a
         # function, as a 'state' that can be updated at will.
 
-        rng = np.random.RandomState(1827)
-        w_init = rng.rand(5)
+        rng = np.random.default_rng(1827)
+        w_init = rng.random((5))
         w = shared(w_init.copy(), "w")
         reg = aet_sum(w * w)
         f = pfunc([], reg)
@@ -127,8 +127,8 @@ def test_param_strict(self):
         out = a + b
 
         f = pfunc([In(a, strict=False)], [out])
-        # works, rand generates float64 by default
-        f(np.random.rand(8))
+        # works, random( generates float64 by default
+        f(np.random.random((8)))
         # works, casting is allowed
         f(np.array([1, 2, 3, 4], dtype="int32"))
 
@@ -145,14 +145,14 @@ def test_param_mutable(self):
 
         # using mutable=True will let fip change the value in aval
         fip = pfunc([In(a, mutable=True)], [a_out], mode="FAST_RUN")
-        aval = np.random.rand(10)
+        aval = np.random.random((10))
         aval2 = aval.copy()
         assert np.all(fip(aval) == (aval2 * 2))
         assert not np.all(aval == aval2)
 
         # using mutable=False should leave the input untouched
         f = pfunc([In(a, mutable=False)], [a_out], mode="FAST_RUN")
-        aval = np.random.rand(10)
+        aval = np.random.random((10))
         aval2 = aval.copy()
         assert np.all(f(aval) == (aval2 * 2))
         assert np.all(aval == aval2)
@@ -375,7 +375,7 @@ def test_update(self):
 
     def test_update_err_broadcast(self):
         # Test that broadcastable dimensions raise error
-        data = np.random.rand(10, 10).astype("float32")
+        data = np.random.random((10, 10)).astype("float32")
         output_var = shared(name="output", value=data)
 
         # the update_var has type matrix, and the update expression

tests/compile/test_debugmode.py

@@ -736,7 +736,7 @@ def perform(self, node, inp, out):
 
 class TestPreallocatedOutput:
     def setup_method(self):
-        self.rng = np.random.RandomState(seed=utt.fetch_seed())
+        self.rng = np.random.default_rng(seed=utt.fetch_seed())
 
     def test_f_contiguous(self):
         a = fmatrix("a")
@@ -745,8 +745,8 @@ def test_f_contiguous(self):
         # In this test, we do not want z to be an output of the graph.
         out = dot(z, np.eye(7))
 
-        a_val = self.rng.randn(7, 7).astype("float32")
-        b_val = self.rng.randn(7, 7).astype("float32")
+        a_val = self.rng.standard_normal((7, 7)).astype("float32")
+        b_val = self.rng.standard_normal((7, 7)).astype("float32")
 
         # Should work
         mode = DebugMode(check_preallocated_output=["c_contiguous"])
@@ -776,8 +776,8 @@ def test_f_contiguous_out(self):
         b = fmatrix("b")
         out = BrokenCImplementationAdd()(a, b)
 
-        a_val = self.rng.randn(7, 7).astype("float32")
-        b_val = self.rng.randn(7, 7).astype("float32")
+        a_val = self.rng.standard_normal((7, 7)).astype("float32")
+        b_val = self.rng.standard_normal((7, 7)).astype("float32")
 
         # Should work
         mode = DebugMode(check_preallocated_output=["c_contiguous"])
@@ -805,5 +805,5 @@ def test_output_broadcast_tensor(self):
         c, r = VecAsRowAndCol()(v)
         f = function([v], [c, r])
 
-        v_val = self.rng.randn(5).astype("float32")
+        v_val = self.rng.standard_normal((5)).astype("float32")
         f(v_val)

tests/compile/test_misc.py

@@ -54,8 +54,8 @@ def __init__(
 
 
 def test_nnet():
-    rng = np.random.RandomState(1827)
-    data = rng.rand(10, 4)
+    rng = np.random.default_rng(279)
+    data = rng.random((10, 4))
     nnet = NNet(n_input=3, n_hidden=10)
     for epoch in range(3):
         mean_cost = 0
@@ -66,7 +66,8 @@ def test_nnet():
             mean_cost += cost
         mean_cost /= float(len(data))
         # print 'Mean cost at epoch %s: %s' % (epoch, mean_cost)
-    assert abs(mean_cost - 0.20588975452) < 1e-6
+    # Seed based test
+    assert abs(mean_cost - 0.2301901) < 1e-6
     # Just call functions to make sure they do not crash.
     nnet.compute_output(input)
     nnet.output_from_hidden(np.ones(10))

tests/d3viz/models.py

@@ -11,7 +11,7 @@ def __init__(self, nfeatures=100, noutputs=10, nhiddens=50, rng=None):
         if rng is None:
             rng = 0
         if isinstance(rng, int):
-            rng = np.random.RandomState(rng)
+            rng = np.random.default_rng(rng)
         self.rng = rng
         self.nfeatures = nfeatures
         self.noutputs = noutputs

tests/d3viz/test_d3viz.py

@@ -19,7 +19,7 @@
 
 class TestD3Viz:
     def setup_method(self):
-        self.rng = np.random.RandomState(0)
+        self.rng = np.random.default_rng(0)
         self.data_dir = pt.join("data", "test_d3viz")
 
     def check(self, f, reference=None, verbose=False):

tests/d3viz/test_formatting.py

@@ -13,7 +13,7 @@
 
 class TestPyDotFormatter:
     def setup_method(self):
-        self.rng = np.random.RandomState(0)
+        self.rng = np.random.default_rng(0)
 
     def node_counts(self, graph):
         node_types = [node.get_attributes()["node_type"] for node in graph.get_nodes()]

tests/link/test_jax.py

@@ -218,6 +218,8 @@ def test_jax_compile_ops():
 
 
 def test_jax_basic():
+    rng = np.random.default_rng(28494)
+
     x = matrix("x")
     y = matrix("y")
     b = vector("b")
@@ -259,15 +261,23 @@ def test_jax_basic():
     out_fg = FunctionGraph([x], [out])
     compare_jax_and_py(
         out_fg,
-        [(np.eye(10) + np.random.randn(10, 10) * 0.01).astype(config.floatX)],
+        [
+            (np.eye(10) + rng.standard_normal(size=(10, 10)) * 0.01).astype(
+                config.floatX
+            )
+        ],
     )
 
     # not sure why this isn't working yet with lower=False
     out = aet_slinalg.Cholesky(lower=False)(x)
     out_fg = FunctionGraph([x], [out])
     compare_jax_and_py(
         out_fg,
-        [(np.eye(10) + np.random.randn(10, 10) * 0.01).astype(config.floatX)],
+        [
+            (np.eye(10) + rng.standard_normal(size=(10, 10)) * 0.01).astype(
+                config.floatX
+            )
+        ],
     )
 
     out = aet_slinalg.solve(x, b)
@@ -294,7 +304,11 @@ def test_jax_basic():
     out_fg = FunctionGraph([x], [out])
     compare_jax_and_py(
         out_fg,
-        [(np.eye(10) + np.random.randn(10, 10) * 0.01).astype(config.floatX)],
+        [
+            (np.eye(10) + rng.standard_normal(size=(10, 10)) * 0.01).astype(
+                config.floatX
+            )
+        ],
     )
 
 
@@ -405,9 +419,9 @@ def test_jax_eye():
 
 
 def test_jax_basic_multiout():
+    rng = np.random.default_rng(213234)
 
-    np.random.seed(213234)
-    M = np.random.normal(size=(3, 3))
+    M = rng.normal(size=(3, 3))
     X = M.dot(M.T)
 
     x = matrix("x")
@@ -638,7 +652,9 @@ def test_jax_Subtensors_omni():
     reason="Omnistaging cannot be disabled",
 )
 def test_jax_IncSubtensor():
-    x_np = np.random.uniform(-1, 1, size=(3, 4, 5)).astype(config.floatX)
+    rng = np.random.default_rng(213234)
+
+    x_np = rng.uniform(-1, 1, size=(3, 4, 5)).astype(config.floatX)
     x_aet = aet.arange(3 * 4 * 5).reshape((3, 4, 5)).astype(config.floatX)
 
     # "Set" basic indices
@@ -661,7 +677,7 @@ def test_jax_IncSubtensor():
 
     # "Set" advanced indices
     st_aet = aet.as_tensor_variable(
-        np.random.uniform(-1, 1, size=(2, 4, 5)).astype(config.floatX)
+        rng.uniform(-1, 1, size=(2, 4, 5)).astype(config.floatX)
     )
     out_aet = aet_subtensor.set_subtensor(x_aet[np.r_[0, 2]], st_aet)
     assert isinstance(out_aet.owner.op, aet_subtensor.AdvancedIncSubtensor1)
@@ -707,7 +723,7 @@ def test_jax_IncSubtensor():
 
     # "Increment" advanced indices
     st_aet = aet.as_tensor_variable(
-        np.random.uniform(-1, 1, size=(2, 4, 5)).astype(config.floatX)
+        rng.uniform(-1, 1, size=(2, 4, 5)).astype(config.floatX)
     )
     out_aet = aet_subtensor.inc_subtensor(x_aet[np.r_[0, 2]], st_aet)
     assert isinstance(out_aet.owner.op, aet_subtensor.AdvancedIncSubtensor1)
@@ -1202,6 +1218,7 @@ def rng_fn(cls, rng, size):
         compare_jax_and_py(fgraph, [])
 
 
+@pytest.mark.xfail(reason="Generators not yet supported in JAX")
 def test_RandomStream():
     srng = RandomStream(seed=123)
     out = srng.normal() - srng.normal()
@@ -1211,3 +1228,11 @@ def test_RandomStream():
     jax_res_2 = fn()
 
     assert np.array_equal(jax_res_1, jax_res_2)
+
+
+@pytest.mark.xfail(reason="Generators not yet supported in JAX")
+def test_random_generators():
+    rng = shared(np.random.default_rng(123))
+    out = normal(rng=rng)
+    fgraph = FunctionGraph([out.owner.inputs[0]], [out], clone=False)
+    compare_jax_and_py(fgraph, [])