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

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

pytensor/link/pytorch/dispatch/random.py

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


It should return a new state, otherwise the draws will be the same the next time it's evaluated

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)

Ch0ronomato · 2024-12-09T01:39:34Z

pytensor/link/pytorch/dispatch/random.py

+@pytorch_typify.register(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)


fwiw, this depends on the device iiuc: https://open.spotify.com/episode/13oJCmQ2JWbk7t6sLRWlDz?si=1f43e67353284cc7

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.

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

pytensor/link/pytorch/dispatch/random.py

+def pytorch_sample_fn_bernoulli(op, node):
+    def sample_fn(rng, size, dtype, p):
+        gen = rng["pytorch_gen"]
+        sample = torch.bernoulli(torch.broadcast_to(p, size), generator=gen)


Size may be None

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

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

pytensor/link/pytorch/dispatch/random.py

+    def sample_fn(rng, size, dtype, n, p):
+        gen = rng["pytorch_gen"]
+        sample = torch.binomial(
+            torch.broadcast_to(n.to(p.dtype), size),


why are you converting n to the type of p?

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

- Copied generator before sampling from it

Ch0ronomato reviewed Nov 10, 2024

View reviewed changes

ricardoV94 reviewed Nov 11, 2024

View reviewed changes

twaclaw added 2 commits December 8, 2024 11:01

Started implementation of random variables with PyTorch backend.

ff973be

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

1c8dc80

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

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

- Added suport for size None

6176479

- Copied generator before sampling from it

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 •

edited

Loading

ricardoV94 Nov 11, 2024

Ch0ronomato 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

ricardoV94 Dec 9, 2024

ricardoV94 Dec 9, 2024

ricardoV94 Dec 9, 2024

	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

ricardoV94 Nov 11, 2024 • 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

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

twaclaw commented Nov 10, 2024 •

edited by github-actions bot

Loading

codecov bot commented Nov 10, 2024 •

edited

Loading

ricardoV94 Nov 11, 2024 •

edited

Loading