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
volatile	Whether this stage's output must be recomputed on every forward pass (rather than reused from cache), invalidating reuse from this point on.

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 and return all cacheable outputs and effects.
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

volatile

Whether this stage’s output must be recomputed on every forward pass (rather than reused from cache), invalidating reuse from this point on.

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 env.evaluation_counter for per-evaluation stochastic sampling). 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:

StageOutput[TypeVar(OutT)]

abstractmethod expand(batch, env)

Transform input and return all cacheable outputs and effects.

Return type:

StageOutput[TypeVar(OutT)]

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