Ensure semantics of rng immutability

Author: Ch0ronomato

pytensor/link/pytorch/dispatch/random.py

@@ -1,3 +1,4 @@
+from copy import deepcopy
 from functools import singledispatch
 
 import numpy as np
@@ -28,7 +29,8 @@ def torch_funcify_RandomVariable(op: ptr.RandomVariable, node, **kwargs):
     rv_sample = pytorch_sample_fn(op, node=node)
 
     def sample_fn(rng, size, *args):
-        return rv_sample(rng, shape, out_dtype, *args)
+        _rng = deepcopy(rng)
+        return rv_sample(_rng, shape, out_dtype, *args)
 
     return sample_fn
 
@@ -57,11 +59,23 @@ def pytorch_sample_fn_binomial(op, node):
     def sample_fn(rng, size, dtype, n, p):
         gen = rng["pytorch_gen"]
         sample = torch.binomial(
-            torch.broadcast_to(n, size),
-            torch.broadcast_to(p, size),
+            torch.broadcast_to(n, size).to(torch.float32),
+            torch.broadcast_to(p, size).to(torch.float32),
             generator=gen,
         )
         rng["pytorch_gen"] = gen
         return (rng, sample)
 
     return sample_fn
+
+
+@pytorch_sample_fn.register(ptr.UniformRV)
+def pytorch_sample_fn_uniform(op, node):
+    def sample_fn(rng, size, dtype, low, high):
+        gen = rng["pytorch_gen"]
+        sample = torch.FloatTensor(size)
+        sample.uniform_(low.item(), high.item(), generator=gen)
+        rng["pytorch_gen"] = gen
+        return (rng, sample)
+
+    return sample_fn

tests/link/pytorch/test_random.py

@@ -16,15 +16,18 @@ def test_random_updates():
     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))
+    f = function([], x, updates={rng: next_rng}, mode="PYTORCH")
+    draws = np.stack([f() for _ in range(5)])
+    # assert we are getting different values
+    assert draws.sum() < 5 and draws.sum() >= 1
 
+    # assert we didn't mutate the original rng
     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)
+        if not isinstance(original.bit_generator.state[k], np.ndarray)
         else np.array_equal(torch_rand[k], original.bit_generator.state[k])
         for k in keys_to_check
     )