[dynamo] Improve support for itertools.product #159693
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/159693
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 5a4251e with merge base c5ec545 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
ciflow/inductor
module: dynamo
topic: not user facing
rtimpe
added
release notes: dynamo
and removed
topic: not user facing
| args = [random.choice(argtypes) for j in range(random.randrange(5))] | ||
| expected_len = prod(map(len, args)) |
There was a problem hiding this comment.
Just comment the for-loop and add a comment explaining why
There was a problem hiding this comment.
re-enabled this section (with the fullgraph=false CM)
torch/_dynamo/polyfills/itertools.py
Outdated
| @substitute_in_graph(itertools.product, is_embedded_type=True) # type: ignore[arg-type] | ||
| def product(*iterables: Iterable[_T], repeat: int = 1) -> Iterator[tuple[_T, ...]]: | ||
| if repeat < 0: | ||
| raise ValueError("repeat argument cannot be negative") | ||
| pools = [tuple(pool) for pool in iterables] * repeat | ||
|
|
||
| def _product(pools: list[tuple[_T, ...]]) -> Iterator[tuple[_T, ...]]: | ||
| result: list[list[_T]] = [[]] | ||
| for pool in pools: | ||
| result = [x + [y] for x in result for y in pool] | ||
|
|
||
| for prod in result: | ||
| yield tuple(prod) | ||
|
|
||
| return _product(pools) |
There was a problem hiding this comment.
I'm ok with this implementation, but it is slow. `list(product("ab", "cd", repeat=8))" takes 12 seconds to finish on my machine.
There was a problem hiding this comment.
I thought this was slow because it's materializing the full output list, but I think it's actually just from dynamo tracing everything. I've gone back to the previous implementation, which delegates to cypthon's product implementation (and is much faster). It might be worth doing the same for permutations - I'll do some testing tomorrow.
There was a problem hiding this comment.
You can run the test with https://github.com/joerick/pyinstrument to see where the bottleneck is
rtimpe
changed the title
| list(product(*args, **dict(repeat=r)))) | ||
| self.assertEqual(len(list(product(*[range(7)]*6))), 7**6) | ||
| self.assertRaises(TypeError, product, range(6), None) | ||
| # self.assertRaises(TypeError, product, range(6), None) |
|
Starting merge as part of PR stack under #159694 |
2 similar comments
|
Starting merge as part of PR stack under #159694 |
|
Starting merge as part of PR stack under #159694 |
|
Starting merge as part of PR stack under #159694 |
Pull Request resolved: #159694 Approved by: https://github.com/guilhermeleobas ghstack dependencies: #159693
Pull Request resolved: pytorch#159693 Approved by: https://github.com/guilhermeleobas, https://github.com/mlazos
Pull Request resolved: pytorch#159694 Approved by: https://github.com/guilhermeleobas ghstack dependencies: pytorch#159693
5 participants
Stack from ghstack (oldest at bottom):
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames @Lucaskabela @mlazos