Update test_binomial

Author: Ch0ronomato

Diff for: tests/link/pytorch/test_random.py

@@ -50,24 +50,22 @@ def test_random_bernoulli(size, p):
 
 
 @pytest.mark.parametrize(
-    "size,n,p",
+    "size,n,p,update",
     [
-        (None, 10, 0.5),
-        ((1000,), 10, 0.5),
-        ((1000, 4), 10, 0.5),
-        ((1000, 2), np.array([10, 40]), np.array([0.5, 0.3])),
+        ((1000,), 10, 0.5, False),
+        ((1000, 4), 10, 0.5, False),
+        ((1000, 2), np.array([10, 40]), np.array([0.5, 0.3]), True),
     ],
 )
-def test_binomial(n, p, size):
+def test_binomial(size, n, p, update):
     rng = shared(np.random.default_rng(123))
-    g = pt.random.binomial(n, p, size=size, rng=rng)
-    g_fn = function([], g, mode=pytorch_mode)
+    rv = pt.random.binomial(n, p, size=size, rng=rng)
+    next_rng, *_ = rv.owner.inputs
+    g_fn = function(
+        [], rv, mode=pytorch_mode, updates={rng: next_rng} if update else None
+    )
     samples = g_fn()
-    if size:
-        np.testing.assert_allclose(samples.mean(axis=0), n * p, rtol=0.1)
-        np.testing.assert_allclose(
-            samples.std(axis=0), np.sqrt(n * p * (1 - p)), rtol=0.2
-        )
-    else:
-        ...
-        # TODO: define test
+    if not update:
+        np.testing.assert_allclose(samples, g_fn(), rtol=0.1)
+    np.testing.assert_allclose(samples.mean(axis=0), n * p, rtol=0.1)
+    np.testing.assert_allclose(samples.std(axis=0), np.sqrt(n * p * (1 - p)), rtol=0.2)