Ensure semantics of rng immutability

Author: Ch0ronomato

pytensor/link/pytorch/dispatch/random.py

@@ -1,23 +1,23 @@
 from functools import singledispatch
 
-import numpy as np
+import numpy.random
 import torch
-from numpy.random import Generator
 
 import pytensor.tensor.random.basic as ptr
 from pytensor.link.pytorch.dispatch.basic import pytorch_funcify, pytorch_typify
 
 
-@pytorch_typify.register(Generator)
+@pytorch_typify.register(numpy.random.Generator)
 def pytorch_typify_Generator(rng, **kwargs):
     # XXX: Check if there is a better way.
     # Numpy uses PCG64 while Torch uses Mersenne-Twister (https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/CPUGeneratorImpl.cpp)
-    state = rng.__getstate__()
-    rng_copy = np.random.default_rng()
-    rng_copy.bit_generator.state = rng.bit_generator.state
-    seed = torch.from_numpy(rng_copy.integers([2**32]))
-    state["pytorch_gen"] = torch.manual_seed(seed)
-    return state
+    seed = torch.from_numpy(rng.integers([2**32]))
+    return torch.manual_seed(seed)
+
+
+@pytorch_typify.register(torch._C.Generator)
+def pytorch_typify_pass_generator(rng, **kwargs):
+    return rng
 
 
 @pytorch_funcify.register(ptr.RandomVariable)
@@ -28,7 +28,9 @@ 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)
+        new_rng = torch.Generator(device="cpu")
+        new_rng.set_state(rng.get_state().clone())
+        return rv_sample(new_rng, shape, out_dtype, *args)
 
     return sample_fn
 
@@ -43,25 +45,31 @@ def pytorch_sample_fn(op, node):
 
 @pytorch_sample_fn.register(ptr.BernoulliRV)
 def pytorch_sample_fn_bernoulli(op, node):
-    def sample_fn(rng, size, dtype, p):
-        gen = rng["pytorch_gen"]
+    def sample_fn(gen, size, dtype, p):
         sample = torch.bernoulli(torch.broadcast_to(p, size), generator=gen)
-        rng["pytorch_gen"] = gen
-        return (rng, sample)
+        return (gen, sample)
 
     return sample_fn
 
 
 @pytorch_sample_fn.register(ptr.BinomialRV)
 def pytorch_sample_fn_binomial(op, node):
-    def sample_fn(rng, size, dtype, n, p):
-        gen = rng["pytorch_gen"]
+    def sample_fn(gen, size, dtype, n, p):
         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 (gen, sample)
+
+    return sample_fn
+
+
+@pytorch_sample_fn.register(ptr.UniformRV)
+def pytorch_sample_fn_uniform(op, node):
+    def sample_fn(gen, size, dtype, low, high):
+        sample = torch.FloatTensor(size)
+        sample.uniform_(low.item(), high.item(), generator=gen)
+        return (gen, sample)
 
     return sample_fn

tests/link/pytorch/test_random.py

@@ -4,30 +4,31 @@
 import pytensor.tensor as pt
 from pytensor.compile.function import function
 from pytensor.compile.sharedvalue import shared
+from pytensor.link.pytorch.dispatch.basic import pytorch_typify
 from tests.link.pytorch.test_basic import pytorch_mode
 
 
 torch = pytest.importorskip("torch")
 
 
-def test_random_updates():
+@pytest.mark.parametrize("update", [(True), (False)])
+def test_random_updates(update):
     original = np.random.default_rng(seed=123)
+    original_torch = pytorch_typify(original)
     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
-    )
+    f = function([], x, updates={rng: next_rng} if update else None, mode="PYTORCH")
+    draws = np.stack([f() for _ in range(5)])
+    # assert we are getting different values
+    if update:
+        assert draws.sum() < 5 and draws.sum() >= 1
+        # assert we have a new rng
+        rng_value = rng.get_value(borrow=True)  # we can't copy torch generator
+        assert torch.eq(rng_value.get_state(), original_torch.get_state())
+    else:
+        pass
 
 
 @pytest.mark.parametrize(
@@ -68,5 +69,10 @@ def test_binomial(size, n, p, update):
     samples = g_fn()
     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)
+        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:
+        second_samples = g_fn()
+        np.testing.assert_array_equal(second_samples, samples)