Implementation of random variables with PyTorch backend #1075

twaclaw · 2024-11-10T15:15:32Z

Description

Related Issue

Closes #
Related to #

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

📚 Documentation preview 📚: https://pytensor--1075.org.readthedocs.build/en/1075/

codecov · 2024-11-10T15:38:24Z

Codecov Report

Attention: Patch coverage is 82.45614% with 10 lines in your changes missing coverage. Please review.

Project coverage is 82.11%. Comparing base (07bd48d) to head (6176479).
Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/link/pytorch/dispatch/random.py	82.60%	8 Missing ⚠️
pytensor/link/pytorch/dispatch/basic.py	50.00%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1075   +/-   ##
=======================================
  Coverage   82.10%   82.11%           
=======================================
  Files         185      186    +1     
  Lines       48089    48184   +95     
  Branches     8659     8673   +14     
=======================================
+ Hits        39485    39564   +79     
- Misses       6439     6452   +13     
- Partials     2165     2168    +3

Files with missing lines	Coverage Δ
pytensor/link/pytorch/dispatch/__init__.py	`100.00% <100.00%> (ø)`
pytensor/link/pytorch/linker.py	`100.00% <100.00%> (ø)`
pytensor/link/pytorch/dispatch/basic.py	`93.69% <50.00%> (-0.81%)`	⬇️
pytensor/link/pytorch/dispatch/random.py	`82.60% <82.60%> (ø)`

... and 3 files with indirect coverage changes

Ch0ronomato · 2024-11-10T23:18:32Z

pytensor/link/pytorch/dispatch/random.py

+    static_shape = rv.type.shape
+    batch_ndim = op.batch_ndim(node)
+
+    # Try to pass static size directly to JAX


nit: pytorch

Ch0ronomato · 2024-11-10T23:20:32Z

pytensor/link/pytorch/dispatch/random.py

+        # XXX replace
+        state_ = rng["pytorch_state"]
+        gen = torch.Generator().set_state(state_)
+        sample = torch.bernoulli(torch.expand_copy(p, size), generator=gen)


I actually don't mind this approach! Torch has a lot of wrapping and abstraction on top of it's random generation, so if we just keep a little bit of state around it feels a bit simpler.

Ch0ronomato · 2024-11-10T23:22:15Z

pytensor/link/pytorch/linker.py

        thunk_inputs = []
        for n in self.fgraph.inputs:
            sinput = storage_map[n]
+            if isinstance(sinput[0], RandomState | Generator):
+                new_value = pytorch_typify(
+                    sinput[0], dtype=getattr(sinput[0], "dtype", None)


Why is this needed?

ricardoV94 · 2024-11-11T07:44:12Z

pytensor/link/pytorch/dispatch/random.py

+    static_shape = rv.type.shape
+    batch_ndim = op.batch_ndim(node)
+
+    # Try to pass static size directly to JAX


This static size is a JAX limitation that shouldn't exist in PyTorch

pytensor/link/pytorch/dispatch/random.py

ricardoV94 · 2024-11-11T07:45:22Z

pytensor/link/pytorch/dispatch/random.py

+        # XXX replace
+        state_ = rng["pytorch_state"]
+        gen = torch.Generator().set_state(state_)
+        sample = torch.bernoulli(torch.expand_copy(p, size), generator=gen)


Shouldn't it jut broadcast?, why copy?

Suggested change

sample = torch.bernoulli(torch.expand_copy(p, size), generator=gen)

sample = torch.bernoulli(torch.expand_copy(p, size), generator=gen)

pytensor/link/pytorch/dispatch/random.py

ricardoV94 · 2024-12-09T15:31:47Z

pytensor/link/pytorch/dispatch/basic.py

+def pytorch_typify(data, dtype=None, **kwargs):
+    if dtype is None:
+        return data
+    else:
+        return torch.tensor(data, dtype=dtype)


We change this approach. You need to dispatch on the RNG type and decide what to do with it. The base-cass is to raise

ricardoV94 · 2024-12-09T15:32:52Z

pytensor/link/pytorch/dispatch/random.py

+    # 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]))


You have to copy the rng before calling rng.integers we don't want to modify the original one

ricardoV94 · 2024-12-09T15:33:54Z

pytensor/link/pytorch/dispatch/random.py

+    def sample_fn(rng, size, *parameters):
+        return pytorch_sample_fn(op, node=node)(rng, shape, out_dtype, *parameters)
+
+    return sample_fn


call pytorch_sample_fn outside of sample_fn.

pytensor/link/pytorch/dispatch/random.py

ricardoV94 · 2024-12-09T15:35:45Z

pytensor/link/pytorch/dispatch/random.py

+        sample = torch.binomial(
+            torch.broadcast_to(n.to(p.dtype), size),
+            torch.broadcast_to(p, size),
+            generator=gen,
+        )
+        return (gen, sample)


size may be none, in which case you should do: n, p = torch.broacast_arrays(n, p) or whatever it's called

pytensor/link/pytorch/dispatch/random.py

ricardoV94 · 2024-12-09T15:37:13Z

pytensor/link/pytorch/linker.py

@@ -84,9 +86,16 @@ def fn(*inputs, inner_fn=inner_fn):
        return fn

    def create_thunk_inputs(self, storage_map):
+        from pytensor.link.pytorch.dispatch import pytorch_typify


You'll need to copy the logic with SharedVariables in JAX to emmit a warning and use different variables. You can refactor the logic so it's not duplicated

ricardoV94 · 2024-12-09T15:37:56Z

tests/link/pytorch/test_random.py

+                4,
+            ),
+            10,
+            0.5,


If you take some of these trailing commas, pre-commit won't force it to be multi-line, which is very unreadable here

ricardoV94 · 2024-12-09T15:38:25Z

tests/link/pytorch/test_random.py

+    ],
+)
+def test_binomial(n, p, size):
+    rng = shared(np.random.default_rng(123))


We need tests that confirm the original rng was not affected

ricardoV94 · 2024-12-09T15:39:33Z

tests/link/pytorch/test_random.py

+    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()


You should call twice. In this case, because you did not set updates you should get the same draws back. See https://pytensor.readthedocs.io/en/latest/tutorial/prng.html for details

You should also test with updates separately

I updated this to include a test without the update, but I'm not getting the same draws. I'll read through the article and see if I can see why

- Copied generator before sampling from it

ricardoV94 · 2025-03-18T07:34:50Z

pytensor/link/pytorch/dispatch/random.py

+    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)


I don't think we need this monkeypatching on the original state anymore, just work directly with the torch.manual_seed. It's not any better to pretend this is still a valid numpy generator

i guess at this point in the pipeline, it doesn't matter if it is still a torch generator, since we're calling typify

ricardoV94 · 2025-03-18T07:36:08Z

pytensor/link/pytorch/dispatch/random.py

+def torch_funcify_RandomVariable(op: ptr.RandomVariable, node, **kwargs):
+    rv = node.outputs[1]
+    out_dtype = rv.type.dtype
+    shape = rv.type.shape


shape is not guaranteed to be static. Use the size argument passed at runtime? Or add an if/else if this was an optimization

ricardoV94 · 2025-03-18T07:36:44Z

pytensor/link/pytorch/dispatch/random.py

+@pytorch_sample_fn.register(ptr.BernoulliRV)
+def pytorch_sample_fn_bernoulli(op, node):
+    def sample_fn(rng, size, dtype, p):
+        gen = rng["pytorch_gen"]


yeah let's not do this indirectaion just work with rng directly, not with the rng stored inside the dictionary

ricardoV94 · 2025-03-18T07:37:59Z

pytensor/link/pytorch/dispatch/random.py

+    rv_sample = pytorch_sample_fn(op, node=node)
+
+    def sample_fn(rng, size, *args):
+        _rng = deepcopy(rng)


you shouldn't always deepcopy, only when op.inplace=False

ricardoV94 · 2025-03-18T07:38:45Z

Looking nearly ready!

Ch0ronomato reviewed Nov 10, 2024

View reviewed changes

ricardoV94 reviewed Nov 11, 2024

View reviewed changes

pytensor/link/pytorch/dispatch/random.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Nov 11, 2024

View reviewed changes

twaclaw force-pushed the implement_random_vars_pytorch_poc branch from 85d6080 to 1c8dc80 Compare December 8, 2024 12:17

Ch0ronomato reviewed Dec 9, 2024

View reviewed changes

pytensor/link/pytorch/dispatch/random.py Show resolved Hide resolved

ricardoV94 reviewed Dec 9, 2024

View reviewed changes

pytensor/link/pytorch/dispatch/random.py Show resolved Hide resolved

ricardoV94 reviewed Dec 9, 2024

View reviewed changes

pytensor/link/pytorch/dispatch/random.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Dec 9, 2024

View reviewed changes

twiecki marked this pull request as draft December 9, 2024 16:30

twiecki changed the title ~~Started implementation of random variables with PyTorch backend [WIP]~~ Implementation of random variables with PyTorch backend Dec 9, 2024

twaclaw and others added 5 commits February 22, 2025 16:12

Started implementation of random variables with PyTorch backend.

69d62c7

Proposal to infer Torch's generator state from the Numpy one

fcd643a

- Added suport for size None

7eb77df

- Copied generator before sampling from it

test rng updates do not overwrite original rng

f93a20e

Update test_binomial

1f863f5

Ch0ronomato force-pushed the implement_random_vars_pytorch_poc branch from 6176479 to 1f863f5 Compare March 11, 2025 05:25

Ch0ronomato added 2 commits March 10, 2025 22:36

Review comments

ac35367

Restore typeify

288d2c4

ricardoV94 reviewed Mar 18, 2025

View reviewed changes

Ensure semantics of rng immutability

4b4f8d0

Ch0ronomato force-pushed the implement_random_vars_pytorch_poc branch from 9f1416a to 4b4f8d0 Compare March 19, 2025 15:53

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implementation of random variables with PyTorch backend #1075

Implementation of random variables with PyTorch backend #1075

twaclaw commented Nov 10, 2024 •

edited by github-actions bot

Loading

codecov bot commented Nov 10, 2024 •

edited

Loading

Ch0ronomato Nov 10, 2024

Ch0ronomato Nov 10, 2024

Ch0ronomato Nov 10, 2024

ricardoV94 Nov 11, 2024

ricardoV94 Nov 11, 2024

ricardoV94 Dec 9, 2024

ricardoV94 Dec 9, 2024

ricardoV94 Dec 9, 2024

ricardoV94 Dec 9, 2024

ricardoV94 Dec 9, 2024

ricardoV94 Dec 9, 2024

ricardoV94 Dec 9, 2024

ricardoV94 Dec 9, 2024

Ch0ronomato Mar 11, 2025

ricardoV94 Mar 18, 2025 •

edited

Loading

Ch0ronomato Mar 19, 2025

ricardoV94 Mar 18, 2025

ricardoV94 Mar 18, 2025

ricardoV94 Mar 18, 2025

ricardoV94 commented Mar 18, 2025

	sample = torch.bernoulli(torch.expand_copy(p, size), generator=gen)
	sample = torch.bernoulli(torch.expand_copy(p, size), generator=gen)

+,
+                          ),
+,
+.5,

Implementation of random variables with PyTorch backend #1075

Are you sure you want to change the base?

Implementation of random variables with PyTorch backend #1075

Conversation

twaclaw commented Nov 10, 2024 • edited by github-actions bot Loading

Description

Related Issue

Checklist

Type of change

codecov bot commented Nov 10, 2024 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 Mar 18, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 commented Mar 18, 2025

twaclaw commented Nov 10, 2024 •

edited by github-actions bot

Loading

codecov bot commented Nov 10, 2024 •

edited

Loading

ricardoV94 Mar 18, 2025 •

edited

Loading