- Added suport for size None

- Copied generator before sampling from it

Author: twaclaw

pytensor/link/pytorch/dispatch/random.py

@@ -1,5 +1,6 @@
 from functools import singledispatch
 
+import numpy as np
 import torch
 from numpy.random import Generator
 
@@ -12,7 +13,9 @@ 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__()
-    seed = torch.from_numpy(rng.integers([2**32]))
+    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
 
@@ -41,8 +44,12 @@ def pytorch_sample_fn(op, node):
 def pytorch_sample_fn_bernoulli(op, node):
     def sample_fn(rng, size, dtype, p):
         gen = rng["pytorch_gen"]
+        if not size:
+            size = (1,)
+
         sample = torch.bernoulli(torch.broadcast_to(p, size), generator=gen)
-        return (gen, sample)
+        rng["pytorch_gen"] = gen
+        return (rng, sample)
 
     return sample_fn
 
@@ -51,11 +58,15 @@ def sample_fn(rng, size, dtype, p):
 def pytorch_sample_fn_binomial(op, node):
     def sample_fn(rng, size, dtype, n, p):
         gen = rng["pytorch_gen"]
+        if not size:
+            size = (1,)
+
         sample = torch.binomial(
             torch.broadcast_to(n.to(p.dtype), size),
             torch.broadcast_to(p, size),
             generator=gen,
         )
-        return (gen, sample)
+        rng["pytorch_gen"] = gen
+        return (rng, sample)
 
     return sample_fn

tests/link/pytorch/test_random.py

@@ -14,13 +14,8 @@
     "size,p",
     [
         ((1000,), 0.5),
-        (
-            (
-                1000,
-                4,
-            ),
-            0.5,
-        ),
+        (None, 0.5),
+        ((1000, 4), 0.5),
         ((10, 2), np.array([0.5, 0.3])),
         ((1000, 10, 2), np.array([0.5, 0.3])),
     ],
@@ -37,29 +32,22 @@ def test_random_bernoulli(size, p):
 @pytest.mark.parametrize(
     "size,n,p",
     [
+        (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, 4), 10, 0.5),
+        ((1000, 2), np.array([10, 40]), np.array([0.5, 0.3])),
     ],
 )
 def test_binomial(n, p, size):
     rng = shared(np.random.default_rng(123))
     g = pt.random.binomial(n, p, size=size, rng=rng)
     g_fn = function([], g, mode=pytorch_mode)
     samples = g_fn()
-    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.1)
+    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