pytorch - ✅(Solved) Fix Dynamo cannot trace `__next__` on `itertools.count` objects. [1 pull requests, 1 comments, 2 participants]

Official PRs (…)
ON THIS PAGE

Recommended Tools

×6

Utilities matched from this issue’s tags and category — try them while you read without losing context.

GitHub issue graph ai analysis

Paste a GitHub issue URL. We fetch that issue, discover linked issues from bodies/comments/timeline, collect linked pull requests, and produce a structured English report.

The report is written in English Markdown for sharing and archival.

Helpful · Quick feedback

Loading…
GitHub stats
pytorch/pytorch#177196Fetched 2026-04-08 00:21:52
View on GitHub
Comments
1
Participants
2
Timeline
42
Reactions
0
Timeline (top)
mentioned ×15subscribed ×15referenced ×5labeled ×4

Error Message

"""Minimal repro: torch.compile fullgraph fails on next(itertools.count()).

Dynamo cannot trace __next__ on itertools.count objects. This reproduces the error from ac.py:157 where _AC_NEXT_REGION_ID = itertools.count() is used inside a compiled region.

Error: torch._dynamo.exc.Unsupported: Dynamo does not know how to trace method __next__ of class count """

import itertools

import torch import torch.nn as nn

_COUNTER = itertools.count()

class MyModule(nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: region_id = next(_COUNTER) # This line triggers the dynamo error return x * (region_id + 1)

model = MyModule() x = torch.randn(4, 4)

This fails with: torch._dynamo.exc.Unsupported: Dynamo does not know how to

trace method __next__ of class count

compiled = torch.compile(model, fullgraph=True) compiled(x)

Fix Action

Fixed

PR fix notes

PR #177876: dynamo: support next() on itertools.count inputs

Description (problem / solution / changelog)

Fix #177196

  1. Root cause Dynamo wrapped pre-existing itertools.count() objects as generic user-defined objects, so next(counter) fell through to the unsupported __next__ path. The existing iterator side-effect replay also had no way to keep a live count() object synchronized after compiled execution.

  2. Proposed fix Wrap pre-existing itertools.count() instances as CountIteratorVariable, guard them on their current iterator state, and replay consumed next() calls back onto the original iterator after the compiled region. Add regression tests for the external-counter fullgraph case and the new guard.

  3. Why the proposed fix is the right long term fix This keeps itertools.count() on Dynamo's existing iterator path instead of treating it like an opaque user object. That gives Dynamo the same three things it already needs for mutable iterators: a dedicated variable type during tracing, guard-based cache invalidation when iterator state changes, and side-effect replay that preserves the original Python object.

Drafted via Codex, published after manual review by @bobrenjc93

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @kadeng @chauhang @amjames @Lucaskabela @jataylo

Changed files

  • test/dynamo/test_guard_serialization.py (modified, +30/-0)
  • test/dynamo/test_misc.py (modified, +15/-0)
  • torch/_dynamo/guards.py (modified, +37/-0)
  • torch/_dynamo/side_effects.py (modified, +9/-0)
  • torch/_dynamo/utils.py (modified, +24/-0)
  • torch/_dynamo/variables/builder.py (modified, +19/-1)
  • torch/_dynamo/variables/iter.py (modified, +24/-4)

Code Example

"""Minimal repro: torch.compile fullgraph fails on next(itertools.count()).

Dynamo cannot trace `__next__` on `itertools.count` objects.
This reproduces the error from ac.py:157 where `_AC_NEXT_REGION_ID = itertools.count()`
is used inside a compiled region.

Error: torch._dynamo.exc.Unsupported: Dynamo does not know how to trace method
`__next__` of class `count`
"""

import itertools

import torch
import torch.nn as nn

_COUNTER = itertools.count()


class MyModule(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        region_id = next(_COUNTER)  # This line triggers the dynamo error
        return x * (region_id + 1)


model = MyModule()
x = torch.randn(4, 4)

# This fails with: torch._dynamo.exc.Unsupported: Dynamo does not know how to
# trace method `__next__` of class `count`
compiled = torch.compile(model, fullgraph=True)
compiled(x)
RAW_BUFFERClick to expand / collapse

🐛 Describe the bug

att

Versions

"""Minimal repro: torch.compile fullgraph fails on next(itertools.count()).

Dynamo cannot trace `__next__` on `itertools.count` objects.
This reproduces the error from ac.py:157 where `_AC_NEXT_REGION_ID = itertools.count()`
is used inside a compiled region.

Error: torch._dynamo.exc.Unsupported: Dynamo does not know how to trace method
`__next__` of class `count`
"""

import itertools

import torch
import torch.nn as nn

_COUNTER = itertools.count()


class MyModule(nn.Module):
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        region_id = next(_COUNTER)  # This line triggers the dynamo error
        return x * (region_id + 1)


model = MyModule()
x = torch.randn(4, 4)

# This fails with: torch._dynamo.exc.Unsupported: Dynamo does not know how to
# trace method `__next__` of class `count`
compiled = torch.compile(model, fullgraph=True)
compiled(x)

cc @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @kadeng @amjames @Lucaskabela @jataylo

extent analysis

Fix Plan

Use a different counter object

The issue arises from using itertools.count() inside a compiled region, which Dynamo cannot trace. We can use a different counter object that Dynamo can handle.

Step-by-Step Solution

  1. Replace itertools.count() with a simple counter variable:

_counter = 0

2. **Update the `forward()` method to use the counter variable**:
   ```python
def forward(self, x: torch.Tensor) -> torch.Tensor:
    region_id = _counter
    _counter += 1  # increment the counter for the next region
    return x * (region_id + 1)
  1. Update the MyModule initialization to reset the counter:

def init(self): super().init() global _counter # access the global counter variable _counter = 0


#### Verification

1. Run the code with the updated counter variable.
2. Check that the code runs without errors.

#### Extra Tips

* Make sure to reset the counter variable in the `__init__` method to ensure it starts from 0 for each new instance of `MyModule`.
* If you need to use a counter object that can be reset, consider using a different data structure, such as a list or a deque, that Dynamo can handle.

Vote matrix · Quick signals

Works
Did the solution work? Tap to confirm.
Easy Fix
Was it a quick fix?
Time Saver
Did it save you time?
Blocking
Was it severely blocking?
Common Issue
Are others likely hitting this too?
Flaky / Intermittent
Is it intermittent?
Verified / Reproducible
Can you reproduce it reliably?
Loading…

Still need to ship something?

×6

Another batch ranked right after the header list — different links, same matching logic.

Back to top recommendations

TRENDING

pytorch - ✅(Solved) Fix Dynamo cannot trace `__next__` on `itertools.count` objects. [1 pull requests, 1 comments, 2 participants]