test rng updates do not overwrite original rng

Author: Ch0ronomato

Diff for: pytensor/link/pytorch/dispatch/random.py

@@ -25,9 +25,10 @@ def torch_funcify_RandomVariable(op: ptr.RandomVariable, node, **kwargs):
     rv = node.outputs[1]
     out_dtype = rv.type.dtype
     shape = rv.type.shape
+    rv_sample = pytorch_sample_fn(op, node=node)
 
-    def sample_fn(rng, size, *parameters):
-        return pytorch_sample_fn(op, node=node)(rng, shape, out_dtype, *parameters)
+    def sample_fn(rng, size, *args):
+        return rv_sample(rng, shape, out_dtype, *args)
 
     return sample_fn
 

Diff for: pytensor/link/pytorch/linker.py

@@ -74,7 +74,9 @@ def __call__(self, *inputs, **kwargs):
                     if getattr(pytensor.link.utils, n[1:], False):
                         delattr(pytensor.link.utils, n[1:])
 
-                return tuple(out.cpu().numpy() for out in outs)
+                return tuple(
+                    out.cpu().numpy() if torch.is_tensor(out) else out for out in outs
+                )
 
             def __del__(self):
                 del self.gen_functors

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

@@ -10,6 +10,26 @@
 torch = pytest.importorskip("torch")
 
 
+def test_random_updates():
+    original = np.random.default_rng(seed=123)
+    rng = shared(original, name="rng", borrow=False)
+    rv = pt.random.bernoulli(0.5, name="y", rng=rng)
+    next_rng, x = rv.owner.outputs
+    x.dprint()
+    f = function([], [x], updates={rng: next_rng}, mode="PYTORCH")
+    assert any(f() for _ in range(5))
+
+    keys_to_check = set(rng.get_value().keys()) - {"pytorch_gen"}
+    torch_rand = {k: rng.get_value()[k] for k in keys_to_check}
+
+    assert all(
+        torch_rand[k] == original.bit_generator.state[k]
+        if not isinstance(original.bit_generator.state, np.ndarray)
+        else np.array_equal(torch_rand[k], original.bit_generator.state[k])
+        for k in keys_to_check
+    )
+
+
 @pytest.mark.parametrize(
     "size,p",
     [