Stage

class Stage(name)

Bases: ABC, Generic[InT, OutT]

Abstract base for pipeline stages.

Attributes Summary

axis_name	Axis name introduced by this stage.
name
stateful	Whether this stage invalidates forward-pass reuse from this point.

Methods Summary

cache_key_extras(env)	Hashable env inputs that this stage reads during expand().
dry_expand(batch, env)	Analytic forward pass for dry runs.
expand(batch, env)	Transform input for the forward pass and return a reduction token.
introspect(batch, env, token)	Return stage-specific metadata for dry-run reporting.
reduce(results, env, token)	Transform results in the backward pass using the forward-pass token.
validate(before, after)	Check this stage's position in the pipeline.

Attributes Documentation

axis_name

Axis name introduced by this stage.

name

stateful

Whether this stage invalidates forward-pass reuse from this point.

Methods Documentation

cache_key_extras(env)

Hashable env inputs that this stage reads during expand().

The forward-pass cache is keyed on the initial spec, the list of stage instances, and the tuple returned by each stage’s cache_key_extras. Override this method to declare any live env state your stage reads during expand (for example env.backend.shots when shot allocation depends on the budget, or a timestamp when the stage depends on external context). Values not listed here won’t trigger cache invalidation when they change. The return value must be hashable; defaults to an empty tuple.

Return type:

tuple[Hashable, ...]

dry_expand(batch, env)

Analytic forward pass for dry runs.

Must emit a batch with the same shape as expand() (same keys, same len(circuit_bodies), same len(measurement_qasms)) and an introspect-compatible token, but without generating expensive per-item content (DAG deep-copies, QASM strings, classical simulations).

The default implementation falls back to expand(), so stages that cannot skip circuit generation keep working unchanged.

Shared-reference contract. Overriding implementations commonly emit the same DAG object across multiple tagged entries in meta.circuit_bodies (true of every built-in dry path today). Consumers may read those DAGs freely, but must not mutate them in place — a mutation on one entry would leak across sibling branches. CircuitPipeline enforces this by demoting such a stage to its real expand() whenever a downstream stage declares consumes_dag_bodies=True and has not overridden dry_expand; a DiviPerformanceWarning is emitted naming both stages. The dry-run circuit count stays correct either way — only the analytic speedup is forfeited for the affected stage.

Return type:

tuple[TypeVar(OutT), Any]

abstractmethod expand(batch, env)

Transform input for the forward pass and return a reduction token.

Return type:

tuple[TypeVar(OutT), Any]

introspect(batch, env, token)

Return stage-specific metadata for dry-run reporting.

Override in subclasses to provide richer introspection data. Called by the dry-run tool after expand with the post-expand batch, the pipeline env, and the stage’s token.

Return type:

dict[str, Any]

abstractmethod reduce(results, env, token)

Transform results in the backward pass using the forward-pass token.

Return type:

dict[Any, Any]

validate(before, after)

Check this stage’s position in the pipeline.

Called by CircuitPipeline at construction time after structural validation. Override to inspect neighboring stages and either:

• raise ContractViolation if preconditions are not met, or

• emit DiviPerformanceWarning for legal-but-slow configurations (e.g. expensive internal options, known-bad neighboring stages). Suppressed at the pipeline level via CircuitPipeline(..., suppress_performance_warnings=True).

Parameters:

• before (tuple[Stage, ...]) – Stages before this one in expand order.

• after (tuple[Stage, ...]) – Stages after this one in expand order.

Return type:

None