Implement ScalarLoop in torch backend #958

Ch0ronomato · 2024-08-01T03:39:08Z

Description

Adds ScalarLoop and Elemwise(ScalarLoop(...)) support for pytorch. Due to #1031 , we have to have a fall back method for things that do what torch calls "data-dependent logic".

Related Issue

Closes Pytorch vmap limitation #1031
Related to Adding conditionals for torch #939

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

pytensor/link/pytorch/dispatch/scalar.py

ricardoV94 · 2024-08-03T14:30:42Z

@Ch0ronomato thanks for taking a stab, I left some comments above

pytensor/link/pytorch/dispatch/scalar.py

codecov · 2024-08-11T19:13:36Z

Codecov Report

Attention: Patch coverage is 81.25000% with 9 lines in your changes missing coverage. Please review.

Project coverage is 82.10%. Comparing base (ef97287) to head (521ad67).
Report is 131 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/link/pytorch/dispatch/scalar.py	77.27%	4 Missing and 1 partial ⚠️
pytensor/link/pytorch/dispatch/elemwise.py	81.81%	2 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #958      +/-   ##
==========================================
+ Coverage   82.09%   82.10%   +0.01%     
==========================================
  Files         183      185       +2     
  Lines       48010    48130     +120     
  Branches     8653     8669      +16     
==========================================
+ Hits        39412    39519     +107     
- Misses       6435     6444       +9     
- Partials     2163     2167       +4

Files with missing lines	Coverage Δ
pytensor/link/pytorch/linker.py	`100.00% <100.00%> (ø)`
pytensor/link/pytorch/dispatch/elemwise.py	`69.11% <81.81%> (+2.45%)`	⬆️
pytensor/link/pytorch/dispatch/scalar.py	`74.07% <77.27%> (+2.19%)`	⬆️

... and 6 files with indirect coverage changes

🚀 New features to boost your workflow:

❄ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

pytensor/link/pytorch/dispatch/scalar.py

tests/link/pytorch/test_basic.py

ricardoV94 · 2024-09-15T08:36:09Z

pytensor/link/pytorch/dispatch/scalar.py

+                carry = update(*carry, *constants)
+            return torch.stack(carry)
+
+    return torch.compiler.disable(scalar_loop)


Can you do recursive=False?

pytensor/link/pytorch/dispatch/scalar.py

Ch0ronomato · 2024-09-19T22:20:25Z

@ricardoV94 - these failures in the CI look a bit strange; i'll look into them before merging...hopefully they go away with merging main 😓

Ch0ronomato · 2024-10-20T23:20:17Z

@ricardoV94 #1031 is blocking the elemwise test - how do you want to proceed with this pr?

ricardoV94 · 2024-10-21T06:33:20Z

@ricardoV94 #1031 is blocking the elemwise test - how do you want to proceed with this pr?

If we can't elemwise it there's not much point to the ScalarLoop. Maybe we need to loop manually instead of vmap for this Op

ricardoV94

I suspect it's in the right direction, but need a bit more help to understand the new code if you can provide it :)

pytensor/link/pytorch/dispatch/elemwise.py

tests/link/pytorch/test_basic.py

Ch0ronomato · 2024-11-03T01:00:15Z

tests/link/pytorch/test_basic.py

+torch_elemwise = pytest.importorskip("pytensor.link.pytorch.dispatch.elemwise")
+
+
+@pytest.mark.parametrize("input_shapes", [[(5, 1, 1, 8), (3, 1, 1), (8,)]])


I set this up so we can try different shapes, but I stuck this one to get started. If you think we should add more lmk.

Ch0ronomato · 2024-11-03T01:01:04Z

tests/link/pytorch/test_basic.py

+    np.testing.assert_equal(mock_inner_func.f.call_count, len(result[0]))
+
+    expected_args = torch.FloatTensor([1.0] * (len(input_shapes) - 1) + [0.0]).unbind(0)
+    expected_calls = starmap(call, repeat(expected_args, mock_inner_func.f.call_count))


I'm bullish on itertools stuff but I think I saw mention earlier that list comprehensions are preferred. I can refactor it if so.

Ch0ronomato · 2024-11-03T01:02:06Z

pytensor/link/pytorch/linker.py

+        from torch import is_tensor
+
+        if is_tensor(out):
+            return out.cpu()


This will probably create conflict when one of my other PRs gets merged as an FYI.

Ch0ronomato · 2024-11-03T01:02:27Z

pytensor/link/pytorch/dispatch/elemwise.py

+                final_inputs[i] = list(layer)
+
+        # make sure we still have the same number of things
+        assert len(final_inputs) == len(shaped_inputs)


I can put these into the unit test if that's preferred now.

If the assert is executed every time at runtime, yes let's not do it here

Ch0ronomato · 2024-11-03T21:26:51Z

tests/link/pytorch/test_basic.py

+        torch.zeros(*input_shapes[-1])
+    ]
+    mock_inner_func = MagicMock()
+    ret_value = torch.rand(2, 2).unbind(0)


Maybe rename to expected

Ch0ronomato · 2024-11-03T21:27:20Z

tests/link/pytorch/test_basic.py

+    mock_inner_func.f.return_value = ret_value
+    elemwise_fn = torch_elemwise.elemwise_scalar_loop(mock_inner_func.f, None, None)
+    result = elemwise_fn(*args)
+    for actual, expected in zip(ret_value, result):


These are backwards fyi

ricardoV94 · 2024-11-12T14:27:44Z

pytensor/link/pytorch/dispatch/elemwise.py

+def elemwise_scalar_loop(base_fn, op, node, **kwargs):
+    """
+    ScalarLoop + Elemwise is too common
+    to not work, but @1031, vmap won't allow it.


Include full link instead of @1031

ricardoV94 · 2024-11-12T14:30:03Z

pytensor/link/pytorch/dispatch/elemwise.py

+        Elemwise._check_runtime_broadcast(node, inputs)
+        shaped_inputs = torch.broadcast_tensors(*inputs)
+        expected_size = shaped_inputs[0].numel()
+        final_inputs = [s.clone() for s in shaped_inputs]


Why .clone()?

This might be unnecessary now. We need the original number of dimensions for the outer loop. I could just grab that count instead.

ricardoV94 · 2024-11-12T14:31:29Z

pytensor/link/pytorch/dispatch/elemwise.py

+        for _ in range(shaped_inputs[0].dim() - 1):
+            for i, _ in enumerate(shaped_inputs):
+                layer = chain.from_iterable([s.unbind(0) for s in final_inputs[i]])
+                final_inputs[i] = list(layer)


What is more performant? Doing this nesting, or raveling all the inputs after broadcasting and doing a single unbind loop?

Either way, doesn't avoid the explicit broadcasting copy or does it?

Ahhhhh, this is basically like ravel you're right!

According to the torch docs, ravel only copies if needed. So there maybe cases where no coping happens

ricardoV94 · 2024-11-12T14:33:42Z

pytensor/link/pytorch/dispatch/elemwise.py

+        assert all(len(x.shape) == 0 for tensor in final_inputs for x in tensor)
+        res = [base_fn(*args) for args in zip(*final_inputs)]
+
+        return [torch.stack(tuple(out[i] for out in res)) for i in range(len(res[0]))]


Will this reintroduce the original shape? Say if the Elemwise of the Scalar Loop had output shape == (5, 3, 2) ?

Co-authored-by: Ricardo Vieira <[email protected]>

ricardoV94 · 2024-12-04T10:40:34Z

pytensor/link/pytorch/dispatch/elemwise.py

+
+        out_shape = bcasted_inputs[0].size()
+        out_size = out_shape.numel()
+        raveled_outputs = [torch.zeros(out_size) for out in node.outputs]


Is there no torch.empty?

mb; i had an old version of torch on my machine (2.2) which didn't have it, but 2.3+ does. Reverted to torch.empty

ricardoV94 · 2024-12-04T10:41:17Z

pytensor/link/pytorch/linker.py

@@ -77,11 +90,11 @@ def __del__(self):
        self.gen_functors = []

        # Torch does not accept numpy inputs and may return GPU objects
-        def fn(*inputs, inner_fn=inner_fn):
+        def create_outputs(*inputs, inner_fn=inner_fn):


Why the new name? Seems less clear

Yeah this fn was shadowing a local variable fn so i just renamed one of them.

Sure but can we use a different name. This doesn't "create_outputs" it converts inputs to torch tensors and outputs back to pytensor-compatible types

sure thing - I can also just keep the shadowing lol. It's not the end of the world.

From your description I would probably have called it convert_types or smth.

We can also put it inside the wrapper __call__ I guess?

Sure, that makes sense too.

pytensor/link/pytorch/linker.py

Ch0ronomato commented Aug 1, 2024

View reviewed changes

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

Ch0ronomato commented Aug 1, 2024

View reviewed changes

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

ricardoV94 requested changes Aug 3, 2024

View reviewed changes

ricardoV94 added enhancement New feature or request torch PyTorch backend labels Aug 3, 2024

Ch0ronomato requested a review from ricardoV94 August 11, 2024 18:56

Ch0ronomato commented Aug 11, 2024

View reviewed changes

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

Ch0ronomato commented Aug 12, 2024

View reviewed changes

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

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

ricardoV94 changed the title ~~Add torch scalar loop~~ Implement ScalarLoop in torch backend Sep 1, 2024

ricardoV94 reviewed Sep 1, 2024

View reviewed changes

tests/link/pytorch/test_basic.py Show resolved Hide resolved

ricardoV94 reviewed Sep 15, 2024

View reviewed changes

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

Ch0ronomato requested a review from ricardoV94 September 19, 2024 22:19

ricardoV94 mentioned this pull request Oct 3, 2024

Implement all Ops in PyTorch (help welcome!) #821

Open

48 tasks

ricardoV94 reviewed Oct 29, 2024

View reviewed changes

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

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

tests/link/pytorch/test_basic.py Outdated Show resolved Hide resolved

Ch0ronomato commented Nov 3, 2024

View reviewed changes

Ch0ronomato requested a review from ricardoV94 November 3, 2024 01:01

Ch0ronomato commented Nov 3, 2024

View reviewed changes

ricardoV94 reviewed Nov 12, 2024

View reviewed changes

Ian Schweer and others added 18 commits November 28, 2024 17:03

Disable compile here based on CI result

2844bc4

Fix mypy signature

39ff3de

Remove unnecessary torch stack

714759c

Only call .cpu when necessary

623dfbe

Recursive false for torch compiler

e06994f

Add elemwise test

977d98d

Late import torch

07e4520

Do iteration instead of vmap for elemwise

561301b

Clean up and add description

e28c3e2

Add unit test to verify iteration

fb90500

Refactor to ravel method

cd678ef

Fix unpacking

1865de9

Co-authored-by: Ricardo Vieira <[email protected]>

Fix comment

1ffd7c6

Remove extra return

fd2f192

Update test

7027c4c

Add single carry test

2a9ffd3

Remove compiler disable

4ebdd15

Better name

46e3e72

Ch0ronomato force-pushed the scalarloop branch from 0905bec to 46e3e72 Compare November 29, 2024 01:05

Ian Schweer added 3 commits November 28, 2024 17:07

Lint

d00f9e2

Better docstring

5bd100e

Fix linker

b6ce485

ricardoV94 reviewed Dec 4, 2024

View reviewed changes

pytensor/link/pytorch/linker.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Dec 4, 2024

View reviewed changes

pytensor/link/pytorch/linker.py Outdated Show resolved Hide resolved

Pr comments

521ad67

Ch0ronomato force-pushed the scalarloop branch from 2f70694 to 521ad67 Compare December 8, 2024 17:00

Ch0ronomato merged commit 9858b33 into pymc-devs:main Dec 8, 2024
60 of 62 checks passed

Ch0ronomato deleted the scalarloop branch December 8, 2024 17:37

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement ScalarLoop in torch backend #958

Implement ScalarLoop in torch backend #958

Ch0ronomato commented Aug 1, 2024 •

edited

Loading

ricardoV94 commented Aug 3, 2024

codecov bot commented Aug 11, 2024 •

edited

Loading

ricardoV94 Sep 15, 2024

Ch0ronomato commented Sep 19, 2024

Ch0ronomato commented Oct 20, 2024

ricardoV94 commented Oct 21, 2024

ricardoV94 left a comment

Ch0ronomato Nov 3, 2024

Ch0ronomato Nov 3, 2024

Ch0ronomato Nov 3, 2024

Ch0ronomato Nov 3, 2024

ricardoV94 Nov 12, 2024 •

edited

Loading

Ch0ronomato Nov 3, 2024

Ch0ronomato Nov 3, 2024

ricardoV94 Nov 12, 2024

ricardoV94 Nov 12, 2024

Ch0ronomato Nov 12, 2024

ricardoV94 Nov 12, 2024 •

edited

Loading

Ch0ronomato Nov 12, 2024

ricardoV94 Nov 12, 2024

ricardoV94 Dec 4, 2024

Ch0ronomato Dec 8, 2024

ricardoV94 Dec 4, 2024

Ch0ronomato Dec 8, 2024

ricardoV94 Dec 8, 2024

Ch0ronomato Dec 8, 2024

ricardoV94 Dec 8, 2024 •

edited

Loading

Ch0ronomato Dec 8, 2024

		torch_elemwise = pytest.importorskip("pytensor.link.pytorch.dispatch.elemwise")


		@pytest.mark.parametrize("input_shapes", [[(5, 1, 1, 8), (3, 1, 1), (8,)]])

Implement ScalarLoop in torch backend #958

Implement ScalarLoop in torch backend #958

Conversation

Ch0ronomato commented Aug 1, 2024 • edited Loading

Description

Related Issue

Checklist

Type of change

ricardoV94 commented Aug 3, 2024

codecov bot commented Aug 11, 2024 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Ch0ronomato commented Sep 19, 2024

Ch0ronomato commented Oct 20, 2024

ricardoV94 commented Oct 21, 2024

ricardoV94 left a 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 Nov 12, 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

ricardoV94 Nov 12, 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

ricardoV94 Dec 8, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Ch0ronomato commented Aug 1, 2024 •

edited

Loading

codecov bot commented Aug 11, 2024 •

edited

Loading

ricardoV94 Nov 12, 2024 •

edited

Loading

ricardoV94 Nov 12, 2024 •

edited

Loading

ricardoV94 Dec 8, 2024 •

edited

Loading