Implementation of scan #134102

bohnstingl · 2024-08-21T11:35:02Z

This operation is supposed to be the pendant to the associative_scan, but can operate with non-associative functions.

cc @albanD @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @chauhang @rec

pytorch-bot · 2024-08-21T11:35:05Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/134102

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 2d99b84 with merge base 32f3af7 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

Reverted while_loop for scan

ydwu4

Please take a look at the comments. Overall it's a bit overwhelming for me to review the diff. Let's clean it up according to the comments and let me know if need another round of reviewing.

torch/_dynamo/variables/higher_order_ops.py

torch/_higher_order_ops/scan.py

torch/_dynamo/variables/higher_order_ops.py

bohnstingl · 2024-08-31T16:24:33Z

@ydwu4, I incorporated your suggestions and updated the code.
I also merge it with the latest main, which includes the associative_scan feature. So there have been a lot of new test cases added and in all of them, I added the scan as well.

I would be more than happy about another round of review if you can. Thank you verymuch!

bohnstingl · 2024-09-07T05:33:30Z

Thank you @ydwu4 for your reviews and your help on this. I fixed your last comments and merged with the latest main.

@albanD, the current PR has been a joint effort with @ydwu4 and went through several iterations of code reviews. It is rather lengthy, but a large portion of that is because we separated scan from associative_scan on a higher level, which produced quite some code lines. Could you please take a look and maybe approve the PR length as it?

zou3519 · 2024-09-09T19:25:22Z

torch/_higher_order_ops/scan.py

+        [pytree.PyTree, pytree.PyTree], Tuple[pytree.PyTree, pytree.PyTree]
+    ],
+    init: pytree.PyTree,
+    input: pytree.PyTree,


I would recommend making everything after input kwarg-only

zou3519 · 2024-09-09T19:26:07Z

torch/_higher_order_ops/scan.py

+    ],
+    init: pytree.PyTree,
+    input: pytree.PyTree,
+    dim: int = 0,


Is anyone asking for dim != 0? The API might be simpler if we don't have this initially.

nvm, I see that associative_scan has a dim argument so this makes sense to have

zou3519 · 2024-09-09T19:27:42Z

torch/_dynamo/variables/higher_order_ops.py

            sub_args,
            sub_kwargs={},
-            description="scan_combine",
+            description="associativ_scan_combine_fn",


associative_scan_combine_fn

zou3519 · 2024-09-09T19:27:48Z

torch/_dynamo/variables/higher_order_ops.py


        combine_gm = torch.fx.GraphModule(dict(tx.output.nn_modules), combine_graph)
-        combine_fn_name = add_subgraph(tx, "scan_combine", combine_gm)
+        combine_fn_name = add_subgraph(tx, "associatie_scan_combine_fn", combine_gm)


This "associative_scan_combine_fn"

zou3519 · 2024-09-09T19:29:15Z

test/functorch/test_control_flow.py

+            torch.float16,
+            torch.float32,
+            torch.float64,
+            torch.int8,
+            torch.int16,
+            torch.int32,
+            torch.int64,
+            torch.complex64,
+            torch.complex128,


Why test so much? This feels like overkill

Yeah, lets keep the commonly seen types e.g. float32, int8, maybe a torch.complex64?

zou3519

from an API perspective:

please make everything after input kwarg only
consider not using "input" as the name of the argument -- input is a special python keyword.

albanD · 2024-09-09T20:00:39Z

Waiving PR length check since every line is carefully reviewed and cannot be logically broken down

ydwu4 · 2024-09-09T20:10:24Z

Talk to @bohnstingl offiline: we'll try to stash the changes to all associative_scan tests and added it back as a follow-up pr.

Refactored generic_scan implementation to increase code re-use

torch/_higher_order_ops/scan.py

ydwu4 · 2024-09-10T04:44:01Z

@pytorchbot merge

pytorchmergebot · 2024-09-10T04:45:50Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

@ydwu4

This operation is supposed to be the pendant to the `associative_scan`, but can operate with non-associative functions. @ydwu4 Pull Request resolved: pytorch#134102 Approved by: https://github.com/ydwu4

ClashLuke · 2025-01-20T10:11:48Z

torch/_higher_order_ops/scan.py

+
+    """
+    if not callable(combine_fn):
+        raise RuntimeError("Combine_fn must be a callable, but got {combine_fn}")


should this be an f-string?

First implementation of scan

2d7d951

pytorch-bot bot added the module: dynamo label Aug 21, 2024

pytorchbot added the open source label Aug 21, 2024

bohnstingl added 2 commits August 21, 2024 23:52

Added RNN test

0905972

WIP: fixed typo and added scan decomposition

ebb57db

pytorch-bot bot added the module: inductor label Aug 23, 2024

bohnstingl added 7 commits August 23, 2024 17:30

Fixed shape checks scan operator

329ec01

Advanced scan operation

5dd844d

Used backend=eager in scan

94fb54f

Reverted while_loop for scan

Updated compile_cnt testcase with reverse flag

3f26caa

Merge branch 'main' of github.com:pytorch/pytorch into scan_operator

e76241b

WIP: Scan after merge with main

b55daf2

Fixed lint issues and skipped scan tests that require inductor

f71629b

bohnstingl mentioned this pull request Aug 28, 2024

[feature request] torch.scan (also port lax.fori_loop / lax.while_loop / lax.associative_scan and hopefully parallelized associative scans) #50688

Open

Removed concatenate operations from generic_scan

3859389

bohnstingl marked this pull request as ready for review August 28, 2024 13:11

bohnstingl added 2 commits August 28, 2024 22:11

Removed unnecessary reverse flag from dispatch keys

a4f42df

Fixed lintrunner and graph test issue

3798454

janeyx99 requested a review from ydwu4 August 30, 2024 04:45

janeyx99 added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Aug 30, 2024

Extended scan with init parameter to be consistent with jax.lax.scan

28c3c9d

ydwu4 reviewed Aug 30, 2024

View reviewed changes

bohnstingl added 3 commits August 31, 2024 18:15

Merge branch 'main' of github.com:pytorch/pytorch into scan_operator

9b27833

Incorporated suggested changes from review

f6c5a5d

Updated higher_order_ops.py

8bb7053

bohnstingl requested a review from ydwu4 August 31, 2024 16:24

bohnstingl added 2 commits September 1, 2024 08:49

Merge branch 'main' of github.com:pytorch/pytorch into scan_operator

d8ef916

Removed unnecessary code

4af0572

Merge branch 'main' of github.com:pytorch/pytorch into scan_operator

0123022

Fixed docstring

c9acb18

ydwu4 added the ciflow/trunk Trigger trunk jobs on your pull request label Sep 7, 2024

zou3519 reviewed Sep 9, 2024

View reviewed changes

zou3519 requested changes Sep 9, 2024

View reviewed changes

albanD added the skip-pr-sanity-checks label Sep 9, 2024

pytorch-bot bot added the ciflow/inductor label Sep 9, 2024

Fixed typos

6dcfba9

Renamed input -> xs for scan and associative_scan

84ed00c

Refactored generic_scan implementation to increase code re-use

ydwu4 reviewed Sep 9, 2024

View reviewed changes

torch/_higher_order_ops/scan.py Show resolved Hide resolved

torch/_higher_order_ops/scan.py Outdated Show resolved Hide resolved

bohnstingl added 2 commits September 10, 2024 00:12

Fixed positional-only arguments for scan

fb0c3a6

Fixed lintrunner issues

2d99b84

pytorchmergebot added the merging label Sep 10, 2024

pytorchmergebot closed this in e889252 Sep 10, 2024

pytorchmergebot added Merged and removed merging labels Sep 10, 2024

ClashLuke reviewed Jan 20, 2025

View reviewed changes

Implementation of scan #134102

Implementation of scan #134102

Uh oh!

Conversation

bohnstingl commented Aug 21, 2024 • edited by pytorch-bot bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Aug 21, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/134102

✅ No Failures

Uh oh!

ydwu4 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

bohnstingl commented Aug 31, 2024

Uh oh!

bohnstingl commented Sep 7, 2024

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ydwu4 Sep 9, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

zou3519 left a comment

Choose a reason for hiding this comment

Uh oh!

albanD commented Sep 9, 2024

Uh oh!

ydwu4 commented Sep 9, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ydwu4 commented Sep 10, 2024

Uh oh!

pytorchmergebot commented Sep 10, 2024

Merge started

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

bohnstingl commented Aug 21, 2024 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented Aug 21, 2024 •

edited

Loading

ydwu4 Sep 9, 2024 •

edited

Loading

ydwu4 commented Sep 9, 2024 •

edited

Loading