iterator

iterator

Iterator-based generation step control.

This module implements tracer.iter[...], which lets the user run a block of intervention code at specific generation steps::

with model.generate("Hello", max_new_tokens=5) as tracer:
    for step in tracer.iter[:]:
        hidden = model.transformer.h[-1].output.save()

Architecture

Each forward pass through the model counts as one generation step. For one-shot intervention hooks to target a specific step, they need to know which step is firing. This is tracked via mediator.iteration_tracker, a defaultdict(int) keyed by provider path (e.g. "model.transformer.h.0.output").

The tracker is not maintained by the intervention hooks themselves — that would require per-module bookkeeping and wouldn't work for modules that are never directly observed. Instead, :class:IteratorTracer registers a set of persistent "tracker-bumping" hooks on every wrapped module when the user enters an iter loop (see :func:register_iter_hooks). These hooks:

• Fire on every forward pass, incrementing the tracker for both the .input and .output paths of the module.
• Use mediator_idx = float('inf') so they run after all intervention and cache hooks — this means the tracker still reflects the "current" step while those hooks check it, and only advances after the step is fully processed.
• Are scoped to the iter loop lifetime: registered at __iter__ entry and removed in the finally block.
• Are per-mediator: each mediator's iter loop installs its own set of hooks that increment only its own tracker.

WrapperModules (generator, streamer, etc.) are skipped because they don't participate in the normal forward-pass cadence — their values are pushed via :meth:eproperty.provide, which bumps the tracker itself through :meth:Interleaver.handle.

IteratorProxy

IteratorProxy(interleaver: Interleaver, model: Envoy)

The object returned by tracer.iter — supports [...] indexing.

Forwards __getitem__ to create an :class:IteratorTracer bound to the current interleaver and the root model envoy. Usage::

for step in tracer.iter[:]:        # all steps
for step in tracer.iter[1:3]:       # steps 1 and 2
for step in tracer.iter[[0, 2, 4]]: # specific steps
for step in tracer.iter[2]:         # single step

interleaver instance-attribute

interleaver = interleaver

model instance-attribute

model = model

__getitem__

__getitem__(iteration: Union[int, slice, list[int]])

IteratorTracer

IteratorTracer(iteration: Union[int, slice, list[int]], interleaver: Interleaver, model: Envoy)

Bases: Tracer

Tracer returned by tracer.iter[...] indexing.

Yields one value per generation step in the requested range and sets mediator.iteration before each yield so that one-shot intervention hooks target the correct forward pass.

On entry, registers persistent tracker-bumping hooks via :func:register_iter_hooks; these are removed in the finally block when the loop completes (normally or via exception).

interleaver instance-attribute

interleaver = interleaver

iteration instance-attribute

iteration = iteration

model instance-attribute

model = model

__iter__

__iter__()

Iterate over the requested generation steps.

For each step i in the range:

1. Set mediator.iteration = i so subsequent intervention requests know which step to target.
2. Yield i to the user's for loop body.
3. On the next loop iteration, repeat.

The original mediator.iteration value is saved on entry and restored in the finally block, so nesting and re-entry work correctly.

compile

compile()

Compile the captured source code as a callable function.

Wraps the captured code in a function definition that accepts the necessary context parameters for execution.

RETURNS	DESCRIPTION
	A callable function that executes the captured code block

execute

execute(fn: Callable)

bump_source_paths

bump_source_paths(mediator, accessor) -> None

Recursively bump iteration trackers for every operation under a SourceAccessor.

Operation-level paths (e.g. "...attn.split_1.output", plus any nested "...attn.X.Y.output" under recursive .source) need their counters advanced in lockstep with the parent module so one-shot operation hooks fire on the right step. The module-level iter hook (see :func:register_iter_hooks) calls this to walk the accessor tree after every forward pass.

The recursion descends through OperationAccessor._source_accessor so deeply nested .source chains stay in sync.

register_iter_hooks

register_iter_hooks(mediator, model) -> List

Register persistent hooks that bump mediator.iteration_tracker for every module after each forward pass.

These hooks are the single source of truth for the per-mediator iteration counter used by :func:hooks.input_hook, :func:hooks.output_hook, and the operation-level equivalents. See the module docstring for the full architecture.

Key properties:

• mediator_idx = float('inf') — fires after all intervention and cache hooks on the same module. This ordering matters: one-shot intervention hooks check tracker[path] == iteration before the iter hook advances the tracker, so the hook is still comparing against the "current" step's value.
• Both .input and .output paths are bumped together from a single output hook. This works because every forward pass runs the input pre-hook chain and the output post-hook chain in lockstep — the two counters stay synchronized.
• When the module has a :class:SourceAccessor, the hook also bumps every operation-level path under it (recursing into nested accessors for recursive .source). The accessor is looked up per fire (via module.__source_accessor__), so an accessor built mid-loop is picked up from its first forward pass onward.
• WrapperModules (generator, streamer, etc.) are skipped. They don't go through PyTorch's forward dispatch on every step; their values flow through :meth:eproperty.provide which bumps the tracker itself.

Known limitation

If a :class:SourceAccessor is built for the first time in step N>0 of the iter loop (i.e. the user's first .source access on that module happens mid-loop), op-path trackers start at 0 instead of N. The user's first hook registered against that accessor captures iteration=N but checks against tracker[op]=0, so that one access misses. Subsequent steps work because per-fire bumping advances both counters together. A future fix may seed op-path trackers from the parent module's tracker at Envoy.source access time.

PARAMETER	DESCRIPTION
mediator	The mediator whose tracker to increment.
model	The root Envoy — the full model.modules() tree is walked to find every wrapped submodule.

RETURNS	DESCRIPTION
List	A list of :class:~torch.utils.hooks.RemovableHandle objects to
List	remove in the iter loop's finally block.