MeasurementStage

class MeasurementStage(grouping_strategy='qwc', result_format_override=None, shot_distribution=None)

Bases: BundleStage

Unified measurement stage for all circuit measurement types.

Handles both probs() and expval(H) measurements, auto-detecting the measurement type from the MetaCircuit. For expval(H), groups observables using the configured strategy (or delegates to expval-native backends via ham_ops).

During expand, sets env.result_format to communicate the expected result format to pipeline.run():

• ResultFormat.PROBS for probs() measurements

• ResultFormat.EXPVALS for expval(H) measurements

• ResultFormat.COUNTS is never set here (reserved for PCE)

During reduce, combines measurement groups using the postprocessing function from _compute_measurement_groups.

Parameters:

• grouping_strategy (Literal['wires', 'default', 'qwc', '_backend_expval'] | None) – Grouping strategy for expval observables. Ignored for probs() measurements. Defaults to “qwc” (qubit-wise commuting).

• result_format_override (ResultFormat | None) – If set, overrides the auto-detected result format. For example, pass ResultFormat.COUNTS to get raw shot counts even when an observable is present.

• shot_distribution (Literal['uniform', 'weighted', 'weighted_random'] | Callable[[list[float], int], list[int]] | None) – How to split the backend’s total shot budget across measurement groups ("uniform", "weighted", "weighted_random", or a callable). Only valid for sampling-based grouping strategies ("qwc", "wires", or None); rejected for "_backend_expval" because the backend computes expectation values analytically and ignores shots. When None (default), every group is submitted with the backend’s full shot count.

Attributes Summary

axis_name	Axis name introduced by this stage.
consumes_dag_bodies	Whether this stage reads meta.circuit_bodies during expand.
handles_measurement	Whether this stage sets up measurement circuits and result format.
stateful	True only for genuinely non-deterministic shot-distribution strategies.

Methods Summary

cache_key_extras(env)	Fold env.backend.shots into the forward-pass cache key.
dry_expand(batch, env)	Analytic path: keep grouping + shot allocation, skip QASM rendering.
expand(batch, env)	Set up measurements and declare the result format on env.
introspect(batch, env, token)	Return stage-specific metadata for dry-run reporting.
reduce(results, env, token)	Combine measurement groups (expval) or strip meas axis (probs).

Attributes Documentation

axis_name

consumes_dag_bodies

handles_measurement

stateful

True only for genuinely non-deterministic shot-distribution strategies.

"weighted_random" and user-supplied callables draw from env.rng and must re-run on every call. The built-in deterministic strategies ("uniform", "weighted") are pure functions of the observable and the shot count, so they stay cacheable; cache invalidation on shot-count changes is delegated to cache_key_extras().

Methods Documentation

cache_key_extras(env)

Fold env.backend.shots into the forward-pass cache key.

Any configured shot distribution (even the deterministic ones) reads env.backend.shots during expand() to compute the per-group allocation. Including the current shot budget in the cache key means a re-run with a different budget triggers fresh allocation rather than replaying the stale one. Returns () when shot_distribution is None.

Return type:

tuple[Hashable, ...]

dry_expand(batch, env)

Analytic path: keep grouping + shot allocation, skip QASM rendering.

Group count is the source of truth for the measurement fan-out, so we always run _compute_measurement_groups and _allocate_per_group_shots (both cheap, pure numpy / analytic). Only the per-group QASM string generation is swapped for a placeholder so the emitted batch has correct shape without materialising measurement circuits.

Return type:

tuple[ExpansionResult, Any]

expand(batch, env)

Set up measurements and declare the result format on env.

Return type:

tuple[ExpansionResult, 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]

reduce(results, env, token)

Combine measurement groups (expval) or strip meas axis (probs).

Return type:

dict[Any, Any]