From ed66358e6f3e262ab9d5d6df579ca0c8e07b59d1 Mon Sep 17 00:00:00 2001 From: Cody Wang Date: Sun, 10 May 2026 14:46:14 -0700 Subject: [PATCH 1/3] feat: Merge program set task results Introduce `merge` method to `ProgramSetQuantumTaskResult` that combines multiple results into one. --- .../tasks/program_set_quantum_task_result.py | 198 +++++++++- .../test_program_set_quantum_task_result.py | 366 ++++++++++++++++++ 2 files changed, 562 insertions(+), 2 deletions(-) diff --git a/src/braket/tasks/program_set_quantum_task_result.py b/src/braket/tasks/program_set_quantum_task_result.py index c4c37b00e..99178171e 100644 --- a/src/braket/tasks/program_set_quantum_task_result.py +++ b/src/braket/tasks/program_set_quantum_task_result.py @@ -16,7 +16,7 @@ import warnings from collections import Counter from collections.abc import Sequence -from dataclasses import dataclass +from dataclasses import dataclass, replace import boto3 import numpy as np @@ -31,10 +31,15 @@ ProgramSetTaskMetadata, ProgramSetTaskResult, ) +from braket.task_result.program_set_executable_result_v1 import ( + ProgramSetExecutableResultMetadata, +) +from braket.task_result.program_set_task_metadata_v1 import ProgramMetadata -from braket.circuits import Observable +from braket.circuits import Circuit, Observable from braket.circuits.observable import EULER_OBSERVABLE_PREFIX from braket.circuits.observables import Sum +from braket.circuits.serialization import IRType from braket.program_sets import CircuitBinding, ParameterSets, ProgramSet from braket.tasks.measurement_utils import ( expectation_from_measurements, @@ -370,6 +375,116 @@ def from_object( program_set=program_set, ) + @staticmethod + def merge( + results: Sequence[ProgramSetQuantumTaskResult], + program_set: ProgramSet, + index_map: list[list[int]], + ) -> ProgramSetQuantumTaskResult: + """Reconstruct a ``ProgramSetQuantumTaskResult`` from the task results produced by running + each program set of ``program_set.split(...)``. + + ``index_map`` is the per-executable map returned alongside the program sets by + ``ProgramSet.split``: ``index_map[k][j]`` gives the index, in the order of ``program_set``, + of the executable that the jth executable of the kth task represents. The kth task's + executables are read in order for its program set, namely across ``results[k].entries``, + and within each ``CompositeEntry`` across its ``entries``. + + The returned ``ProgramSetQuantumTaskResult`` has the same shape as if ``program_set`` had + been run unsplit, namely one ``CompositeEntry`` per entry of ``program_set.entries``, + and ``MeasuredEntry`` objects in the order of the program. + + Expectation values and ``Sum`` Hamiltonian expectations are computed + for the original ``ProgramSet``. + + Args: + results (Sequence[ProgramSetQuantumTaskResult]): The result of each task, in the same + order as ``program_set.split``'s return. + program_set (ProgramSet): The original unsplit program set. + index_map (list[list[int]]): The per-executable map from ``ProgramSet.split``. + + Returns: + ProgramSetQuantumTaskResult: A result matching the shape of ``program_set``. + + Raises: + ValueError: If ``len(results) != len(index_map)``, if the total size of ``index_map`` + doesn't match ``program_set.total_executables``, or if any task produces a + different number of executables than its map expects. + """ + if len(results) != len(index_map): + raise ValueError( + f"Got {len(results)} task results but {len(index_map)} entries in index_map" + ) + total_executables = program_set.total_executables + total_mapped = sum(len(m) for m in index_map) + if total_mapped != total_executables: + raise ValueError( + f"Index map covers {total_mapped} executables but the original program set " + f"has {total_executables}" + ) + + programs = [_binding_to_program(binding) for binding in program_set.entries] + executable_indices = list(program_set.enumerate_executables()) + binding_executable_counts = [_count_executables(b) for b in program_set.entries] + shots_per_executable = results[0].entries[0].shots_per_executable + + buffer = [None] * total_executables + for k, result in enumerate(results): + _buffer_result( + k=k, + result=result, + map_k=index_map[k], + program_set=program_set, + programs=programs, + executable_indices=executable_indices, + buffer=buffer, + ) + + entries = [] + start = 0 + for binding_idx, binding in enumerate(program_set.entries): + count = binding_executable_counts[binding_idx] + program = programs[binding_idx] + observables = binding.observables if isinstance(binding, CircuitBinding) else None + entries.append( + CompositeEntry( + entries=buffer[start : start + count], + program=program, + inputs=CompositeEntry._get_inputs(program, observables), + observables=observables, + shots_per_executable=shots_per_executable, + additional_metadata=None, + ) + ) + start += count + + metas = [r.task_metadata for r in results] + return ProgramSetQuantumTaskResult( + entries=entries, + task_metadata=ProgramSetTaskMetadata( + id=";".join(meta.id for meta in metas), # Better way to do this? + deviceId=metas[0].deviceId, + requestedShots=sum(m.requestedShots for m in metas), + successfulShots=sum(m.successfulShots for m in metas), + programMetadata=[ + ProgramMetadata( + executables=[ + ProgramSetExecutableResultMetadata() + for _ in range(_count_executables(b)) + ] + ) + for b in program_set.entries + ], + deviceParameters=None, # TODO: find a way to fill this in + createdAt=min(m.createdAt for m in metas if m.createdAt), + endedAt=max(m.endedAt for m in metas if m.endedAt), + status="COMPLETED" if any(m.status == "COMPLETED" for m in metas) else "FAILED", + totalFailedExecutables=sum(m.totalFailedExecutables for m in metas), + ), + num_executables=total_executables, + program_set=program_set, + ) + def __len__(self): return len(self.entries) @@ -481,6 +596,85 @@ def _compute_num_executables(metadata: ProgramSetTaskMetadata) -> int: return counter +def _binding_to_program(binding: CircuitBinding | Circuit) -> Program: + if isinstance(binding, Circuit): + return Program(source=binding.to_ir(IRType.OPENQASM).source, inputs=None) + return binding.to_ir() + + +def _count_executables(binding: CircuitBinding | Circuit) -> int: + if isinstance(binding, Circuit): + return 1 + num_ps = len(binding.input_sets) if binding.input_sets is not None else 1 + num_obs = len(binding.observables) if binding.observables is not None else 1 + return num_ps * num_obs + + +def _buffer_result( + k: int, + result: ProgramSetQuantumTaskResult, + map_k: list[int], + program_set: ProgramSet, + programs: list[Program], + executable_indices: list[tuple[int, int, int]], + buffer: list[MeasuredEntry | ProgramSetExecutableFailure | None], +) -> None: + j = 0 + for composite in result.entries: + for entry in composite.entries: + if j >= len(map_k): + raise ValueError( + f"t=Task {result.task_metadata.id} at index {k} " + "produced more executables than index map expects" + ) + orig_idx = map_k[j] + binding_idx, ps_idx, obs_idx = executable_indices[orig_idx] + buffer[orig_idx] = _convert_measured_entry( + entry, + program_set.entries[binding_idx], + programs[binding_idx], + ps_idx, + obs_idx, + ) + j += 1 + if j != len(map_k): + raise ValueError( + f"Task {result.task_metadata.id} at index {k} produced {j} executables " + f"but index map expected {len(map_k)}" + ) + + +def _convert_measured_entry( + entry: MeasuredEntry | ProgramSetExecutableFailure, + original_binding: CircuitBinding | Circuit, + original_program: Program, + parameter_set_index: int, + observable_index: int, +) -> MeasuredEntry | ProgramSetExecutableFailure: + if isinstance(entry, ProgramSetExecutableFailure): + return entry + if isinstance(original_binding, Circuit): + return replace(entry, program=original_program.source, inputs=None, observable=None) + observables = original_binding.observables + if observables is None: + observable: Observable | None = None + num_obs = 1 + elif isinstance(observables, Sum): + observable = observables.summands[observable_index] + num_obs = len(observables.summands) + else: + observable = observables[observable_index] + num_obs = len(observables) + orig_inputs_index = parameter_set_index * num_obs + observable_index + program_inputs = original_program.inputs or {} + return replace( + entry, + program=original_program.source, + inputs={key: value[orig_inputs_index] for key, value in program_inputs.items()} or None, + observable=observable, + ) + + def _retrieve_s3_object_body(s3_bucket: str, s3_object_key: str, s3_client: BaseClient) -> str: """Retrieve the S3 object body. diff --git a/test/unit_tests/braket/tasks/test_program_set_quantum_task_result.py b/test/unit_tests/braket/tasks/test_program_set_quantum_task_result.py index 3580c4eb0..e3170daaf 100644 --- a/test/unit_tests/braket/tasks/test_program_set_quantum_task_result.py +++ b/test/unit_tests/braket/tasks/test_program_set_quantum_task_result.py @@ -19,6 +19,8 @@ from braket.circuits import Circuit from braket.circuits.observables import X, Y, Z +from braket.circuits.serialization import IRType +from braket.ir.openqasm import Program from braket.parametric import FreeParameter from braket.program_sets import CircuitBinding, ParameterSets, ProgramSet from braket.schema_common import BraketSchemaBase @@ -437,3 +439,367 @@ def test_dispatch_executable_result_with_none_inputs(execution_measurement_proba assert isinstance(measured_entry, MeasuredEntry) assert measured_entry.inputs is None assert measured_entry.probabilities == {"00": 0.7, "11": 0.3} + + +_SIM_METADATA_HEADER = { + "braketSchemaHeader": {"name": "braket.task_result.simulator_metadata", "version": "1"}, + "executionDuration": 50, +} +_DEVICE_PARAMS = { + "braketSchemaHeader": { + "name": "braket.device_schema.simulators.gate_model_simulator_device_parameters", + "version": "1", + }, + "paradigmParameters": { + "braketSchemaHeader": { + "name": "braket.device_schema.gate_model_parameters", + "version": "1", + }, + "qubitCount": 5, + "disableQubitRewiring": False, + }, +} + + +def _make_exec_result(inputs_index, probs=None): + return { + "braketSchemaHeader": { + "name": "braket.task_result.program_set_executable_result", + "version": "1", + }, + "inputsIndex": inputs_index, + "measurementProbabilities": probs or {"00": 0.7, "11": 0.3}, + "measuredQubits": [0, 1], + } + + +def _make_program_result(program_dict, executable_dicts): + return { + "braketSchemaHeader": {"name": "braket.task_result.program_result", "version": "1"}, + "executableResults": executable_dicts, + "source": program_dict, + "additionalMetadata": {"simulatorMetadata": dict(_SIM_METADATA_HEADER)}, + } + + +def _make_task_metadata( + program_executable_counts, task_id="arn:aws:braket:::task/sub", shots_per_executable=40 +): + total = sum(program_executable_counts) + return { + "braketSchemaHeader": { + "name": "braket.task_result.program_set_task_metadata", + "version": "1", + }, + "id": task_id, + "deviceId": "arn:aws:braket:::device/quantum-simulator/amazon/sv1", + "requestedShots": shots_per_executable * total, + "successfulShots": shots_per_executable * total, + "programMetadata": [ + {"executables": [{} for _ in range(n)]} for n in program_executable_counts + ], + "deviceParameters": dict(_DEVICE_PARAMS), + "createdAt": "2024-10-15T19:06:58.986Z", + "endedAt": "2024-10-15T19:07:00.382Z", + "status": "COMPLETED", + "totalFailedExecutables": 0, + } + + +def _make_task_result(program_results, metadata): + return { + "braketSchemaHeader": { + "name": "braket.task_result.program_set_task_result", + "version": "1", + }, + "programResults": program_results, + "taskMetadata": metadata, + } + + +def _parse(d): + return BraketSchemaBase.parse_raw_schema(json.dumps(d)) + + +def _build_sub_quantum_result(sub_program_set, programs_execs, shots_per_executable=40): + """Build a :class:`ProgramSetQuantumTaskResult` for a sub-program-set by first + building a wire-format ``ProgramSetTaskResult`` and passing it through + :meth:`ProgramSetQuantumTaskResult.from_object`. + + Args: + sub_program_set: The sub-``ProgramSet`` whose run produced the result. + programs_execs: One list of exec-result dicts per entry in ``sub_program_set.entries``. + shots_per_executable: shots per executable, propagated to the metadata. + """ + program_results = [] + counts = [] + for entry, execs in zip(sub_program_set.entries, programs_execs, strict=True): + if isinstance(entry, CircuitBinding): + source_dict = entry.to_ir().dict() + else: + source_dict = Program(source=entry.to_ir(IRType.OPENQASM).source, inputs=None).dict() + program_results.append(_make_program_result(source_dict, execs)) + counts.append(len(execs)) + wire = _parse( + _make_task_result( + program_results, _make_task_metadata(counts, shots_per_executable=shots_per_executable) + ) + ) + return ProgramSetQuantumTaskResult.from_object(wire, sub_program_set) + + +def test_from_multiple_single_sub_task_no_split_roundtrips(circuit_rx_parametrized_fixture): + """If split returns [self], from_multiple should reproduce from_object's output.""" + binding = CircuitBinding( + circuit_rx_parametrized_fixture, + input_sets={"theta": [0.12, 2.1]}, + observables=10 * Z(0) + X(0) - 0.01 * Y(0) @ X(1), + ) + ps = ProgramSet(binding) + subs, index_map = ps.split(100) # fits, so one sub-task identical to ps. + assert subs == [ps] + + # Build a ProgramSetQuantumTaskResult that represents running this ps: the wire + # payload goes through from_object first. + sub_program = subs[0].to_ir().programs[0].dict() + execs = [_make_exec_result(i) for i in range(ps.total_executables)] + wire = _parse( + _make_task_result( + [_make_program_result(sub_program, execs)], + _make_task_metadata([ps.total_executables]), + ) + ) + reference = ProgramSetQuantumTaskResult.from_object(wire, ps) + + merged = ProgramSetQuantumTaskResult.merge([reference], ps, index_map) + + assert len(merged) == len(reference) == 1 + ref_composite = reference[0] + got_composite = merged[0] + assert len(got_composite) == len(ref_composite) + assert got_composite.program == ref_composite.program + assert got_composite.inputs == ref_composite.inputs + assert got_composite.observables == ref_composite.observables + for m_got, m_ref in zip(got_composite.entries, ref_composite.entries): + assert m_got.measured_qubits == m_ref.measured_qubits + assert m_got.probabilities == m_ref.probabilities + assert m_got.observable == m_ref.observable + assert m_got.inputs == m_ref.inputs + + +def test_from_multiple_split_list_observables(circuit_rx_parametrized_fixture): + """Split a binding with more observables than fit; scatter + regroup must + reconstruct the same CompositeEntry as running unsplit.""" + binding = CircuitBinding( + circuit_rx_parametrized_fixture, + input_sets={"theta": [0.12]}, + observables=[X(0), Y(0), Z(0), X(0) @ Y(1)], # 4 observables. + ) + ps = ProgramSet(binding) + subs, index_map = ps.split(2) # 4 > 2, so observables split into windows (0,2), (2,4). + assert [s.total_executables for s in subs] == [2, 2] + + # One sub-quantum-result per sub-program-set, built by running each through + # from_object on an inline wire payload. + sub_results = [ + _build_sub_quantum_result( + sub, [[_make_exec_result(i, {"00": 1.0}) for i in range(sub.total_executables)]] + ) + for sub in subs + ] + + merged = ProgramSetQuantumTaskResult.merge(sub_results, ps, index_map) + assert len(merged) == 1 + composite = merged[0] + # The merged composite should have 4 MeasuredEntries in canonical order, each with + # the ORIGINAL binding's observable attached at that index. + assert len(composite) == 4 + for i, measured in enumerate(composite.entries): + assert isinstance(measured, MeasuredEntry) + assert measured.observable == binding.observables[i] + assert composite.inputs == ParameterSets({"theta": [0.12]}) + # task metadata was aggregated across sub-tasks. + assert merged.num_executables == 4 + assert merged.task_metadata.requestedShots == sum( + r.task_metadata.requestedShots for r in sub_results + ) + assert merged.task_metadata.successfulShots == sum( + r.task_metadata.successfulShots for r in sub_results + ) + + +def test_from_multiple_split_sum_hamiltonian_reconstructs_expectation( + circuit_rx_parametrized_fixture, +): + """Splitting a Sum Hamiltonian across multiple sub-tasks and then merging must + reconstruct the full expectation value, because scatter+regroup feeds the original + Sum back into ``_compute_expectations``.""" + # Same fixture as existing test_observables_no_inputs (with known expectation). + circuit = Circuit().h(0).cnot(0, 1) + h = 10000 * Z(0) + 1000 * X(0) - 100 * Z(0) + 10 * Z(1) + X(1) - 0.1 * Y(1) + binding = CircuitBinding(circuit, observables=h) + ps = ProgramSet(binding) + assert ps.total_executables == 6 + + subs, index_map = ps.split(2) # 6 > 2, so Sum splits into 3 windows of size 2. + assert [s.total_executables for s in subs] == [2, 2, 2] + + # Each executable's measurement is the same {"00": 0.7, "11": 0.3} as the existing + # test_observables_no_inputs fixture, so the expectation should match 4364.36. + sub_results = [ + _build_sub_quantum_result( + sub, [[_make_exec_result(i) for i in range(sub.total_executables)]] + ) + for sub in subs + ] + + merged = ProgramSetQuantumTaskResult.merge(sub_results, ps, index_map) + composite = merged[0] + assert composite.observables is h + assert len(composite) == 6 + assert np.isclose(composite.expectation(), 4364.36) + + +def test_from_multiple_mixed_bindings_and_failures(circuit_rx_parametrized_fixture): + """A program set with multiple bindings, split across sub-tasks, with one + executable failing in a sub-task. Failures must land at the correct original + position in the merged result.""" + c1 = circuit_rx_parametrized_fixture + c2 = Circuit().rx(0, FreeParameter("phi")) + b1 = CircuitBinding(c1, {"theta": [0.1, 0.2, 0.3]}, observables=[X(0), Y(0)]) # 6 execs + b2 = CircuitBinding(c2, {"phi": [0.4, 0.5]}) # 2 execs, no observables + ps = ProgramSet([b1, b2]) + assert ps.total_executables == 8 + + subs, index_map = ps.split(5) + # Greedy pack with max=5: b1 classes (sizes 2,2,2) fill [2+2=4, +2>5 flush], so + # sub 0 = 2 b1 classes (4 execs), sub 1 = 1 b1 class (2 execs) + b2 (2 execs) = 4 execs. + assert [s.total_executables for s in subs] == [4, 4] + + def _failure(inputs_index): + return { + "braketSchemaHeader": { + "name": "braket.task_result.program_set_executable_failure", + "version": "1", + }, + "inputsIndex": inputs_index, + "failureMetadata": { + "failureReason": "test failure", + "retryable": False, + "category": "DEVICE", + }, + } + + # Inject a failure at original index 5 (b1 ps=2, obs=1) which lives in sub 1. + sub_results = [] + failure_injected = False + for k, sub in enumerate(subs): + programs_execs = [] + for prog_idx, entry in enumerate(sub.entries): + num_execs = len(entry) if isinstance(entry, CircuitBinding) else 1 + execs = [] + for i in range(num_execs): + # Figure out this sub-executable's original index. Within sub k, + # j runs across all programs so we need a running counter. + j = ( + sum( + len(prev_entry) if isinstance(prev_entry, CircuitBinding) else 1 + for prev_entry in sub.entries[:prog_idx] + ) + + i + ) + if index_map[k][j] == 5: + execs.append(_failure(i)) + failure_injected = True + else: + execs.append(_make_exec_result(i)) + programs_execs.append(execs) + sub_results.append(_build_sub_quantum_result(sub, programs_execs)) + + assert failure_injected + merged = ProgramSetQuantumTaskResult.merge(sub_results, ps, index_map) + assert len(merged) == 2 + # Binding 0: 6 executables, position 5 is a failure. + assert len(merged[0]) == 6 + # Binding 1: 2 executables, all successful. + assert len(merged[1]) == 2 + from braket.task_result import ProgramSetExecutableFailure + + assert isinstance(merged[0].entries[5], ProgramSetExecutableFailure) + # All non-failure entries for binding 0 have the correct observables. + for i, entry in enumerate(merged[0].entries): + if isinstance(entry, MeasuredEntry): + expected_obs = b1.observables[i % len(b1.observables)] + assert entry.observable == expected_obs + # Binding 1 entries have no observable. + for entry in merged[1].entries: + if isinstance(entry, MeasuredEntry): + assert entry.observable is None + + +def test_from_multiple_validates_index_map_size(circuit_rx_parametrized_fixture): + binding = CircuitBinding(circuit_rx_parametrized_fixture, input_sets={"theta": [0.1, 0.2]}) + ps = ProgramSet(binding) + sub_result = _build_sub_quantum_result(ps, [[_make_exec_result(0), _make_exec_result(1)]]) + # index_map has 1 entry for 1 sub-task, but size doesn't match ps.total_executables. + with pytest.raises(ValueError, match="Index map covers 1"): + ProgramSetQuantumTaskResult.merge([sub_result], ps, [[0]]) + # Sub-task count doesn't match index_map's length. + with pytest.raises(ValueError, match="1 task results but 2 entries in index_map"): + ProgramSetQuantumTaskResult.merge([sub_result], ps, [[0], [1]]) + + +@pytest.fixture +def circuit_rx_parametrized_fixture(): + return Circuit().rx(0, FreeParameter("theta")).cnot(0, 1) + + +def test_from_multiple_with_plain_circuit_entries(): + """from_multiple should handle plain Circuit entries (no inputs, no observables).""" + c1 = ghz_test(2) + c2 = ghz_test(1) + ps = ProgramSet([c1, c2]) + subs, index_map = ps.split(1) + assert [s.total_executables for s in subs] == [1, 1] + + sub_results = [_build_sub_quantum_result(sub, [[_make_exec_result(0)]]) for sub in subs] + + merged = ProgramSetQuantumTaskResult.merge(sub_results, ps, index_map) + assert len(merged) == 2 + assert len(merged[0]) == 1 + assert len(merged[1]) == 1 + assert merged[0].observables is None + assert merged[0].entries[0].observable is None + assert merged[0].entries[0].inputs is None + + +def test_from_multiple_rejects_sub_task_over_index_map(circuit_rx_parametrized_fixture): + """Sub-task has more executables than index_map[k] covers.""" + binding = CircuitBinding(circuit_rx_parametrized_fixture, input_sets={"theta": [0.1, 0.2]}) + ps = ProgramSet(binding) + # Sub-task reports 2 executables, but index_map says there's only 1. + sub_result = _build_sub_quantum_result(ps, [[_make_exec_result(0), _make_exec_result(1)]]) + with pytest.raises(ValueError, match="produced more executables than index map"): + ProgramSetQuantumTaskResult.merge( + [sub_result], + ProgramSet(CircuitBinding(circuit_rx_parametrized_fixture, {"theta": [0.1]})), + [[0]], + ) + + +def test_from_multiple_rejects_sub_task_under_index_map(circuit_rx_parametrized_fixture): + """Sub-task has fewer executables than index_map[k] covers.""" + binding = CircuitBinding(circuit_rx_parametrized_fixture, input_sets={"theta": [0.1, 0.2]}) + ps = ProgramSet(binding) + # Sub-task reports only 1 executable, but index_map says there are 2. + sub_result = _build_sub_quantum_result(ps, [[_make_exec_result(0)]]) + with pytest.raises(ValueError, match="expected 2"): + ProgramSetQuantumTaskResult.merge([sub_result], ps, [[0, 1]]) + + +def ghz_test(n): + """Local ghz helper so tests don't depend on program_set_test_utils.""" + circuit = Circuit().h(0) + for i in range(n - 1): + circuit.cnot(i, i + 1) + return circuit From 726bcee099852515f0cb4e68fd9a0720a9f8c997 Mon Sep 17 00:00:00 2001 From: Cody Wang Date: Wed, 27 May 2026 14:29:07 -0700 Subject: [PATCH 2/3] Add warning for merged metadata And other minor changes --- .../tasks/program_set_quantum_task_result.py | 160 +++++++++++++----- .../test_program_set_quantum_task_result.py | 125 ++++++++++++-- 2 files changed, 233 insertions(+), 52 deletions(-) diff --git a/src/braket/tasks/program_set_quantum_task_result.py b/src/braket/tasks/program_set_quantum_task_result.py index 99178171e..38e8b15c9 100644 --- a/src/braket/tasks/program_set_quantum_task_result.py +++ b/src/braket/tasks/program_set_quantum_task_result.py @@ -146,7 +146,6 @@ def expectation(self) -> float | None: return self._expectation -@dataclass class CompositeEntry: """Results of a program in a program set @@ -157,15 +156,70 @@ class CompositeEntry: observables (Sum | list[Observable] | None): The Sum Hamiltonian or observables that were measured, if any. shots_per_executable (int): The number of shots each underlying executable was run with - additional_metadata (AdditionalMetadata): Additional metadata about this program + additional_metadata (AdditionalMetadata | None): Additional metadata about this program. + ``None`` for entries produced by ``ProgramSetQuantumTaskResult.merge``, + since per-program metadata cannot be aggregated meaningfully across underlying tasks. """ - entries: list[MeasuredEntry] - program: Program - inputs: ParameterSets - observables: Sum | list[Observable] | None - shots_per_executable: int - additional_metadata: AdditionalMetadata + def __init__( + self, + entries: list[MeasuredEntry], + program: Program, + inputs: ParameterSets, + observables: Sum | list[Observable] | None, + shots_per_executable: int, + additional_metadata: AdditionalMetadata | None, + ): + self._entries = entries + self._program = program + self._inputs = inputs + self._observables = observables + self._shots_per_executable = shots_per_executable + self._additional_metadata = additional_metadata + self._was_merged = False + self._expectations = self._compute_expectations() if isinstance(observables, Sum) else None + + @property + def entries(self) -> list[MeasuredEntry]: + """list[MeasuredEntry]: The results of each executable in this program.""" + return self._entries + + @property + def program(self) -> Program: + """Program: The program that was run.""" + return self._program + + @property + def inputs(self) -> ParameterSets: + """ParameterSets: The input values this program was run with.""" + return self._inputs + + @property + def observables(self) -> Sum | list[Observable] | None: + """Sum | list[Observable] | None: The Sum Hamiltonian or observables measured, + if any.""" + return self._observables + + @property + def shots_per_executable(self) -> int: + """int: The number of shots each underlying executable was run with.""" + return self._shots_per_executable + + @property + def additional_metadata(self) -> AdditionalMetadata | None: + """AdditionalMetadata | None: Additional metadata about this program. + + For entries produced by ``ProgramSetQuantumTaskResult.merge``, this will be ``None``; + Use the original per-task results for true per-program metadata. + """ + if self._was_merged: + warnings.warn( + "additional_metadata for a CompositeEntry on a merged " + "ProgramSetQuantumTaskResult is None; " + "use the original per-task results for true per-program metadata.", + stacklevel=2, + ) + return self._additional_metadata @staticmethod def _from_object( @@ -197,11 +251,6 @@ def _from_object( additional_metadata=program_result.additionalMetadata, ) - def __post_init__(self): - self._expectations = ( - self._compute_expectations() if isinstance(self.observables, Sum) else None - ) - def __len__(self): return len(self.entries) @@ -315,7 +364,6 @@ def _dispatch_executable_result( ) -@dataclass class ProgramSetQuantumTaskResult: """The result of a program set task. @@ -328,10 +376,45 @@ class ProgramSetQuantumTaskResult: can be automatically computed. """ - entries: list[CompositeEntry] - task_metadata: ProgramSetTaskMetadata - num_executables: int - program_set: ProgramSet | None + def __init__( + self, + entries: list[CompositeEntry], + task_metadata: ProgramSetTaskMetadata, + num_executables: int, + program_set: ProgramSet | None, + ): + self._entries = entries + self._task_metadata = task_metadata + self._num_executables = num_executables + self._program_set = program_set + self._was_merged = False + + @property + def entries(self) -> list[CompositeEntry]: + """list[CompositeEntry]: The results of each program in this program set.""" + return self._entries + + @property + def task_metadata(self) -> ProgramSetTaskMetadata: + """ProgramSetTaskMetadata: The metadata of the task.""" + if self._was_merged: + warnings.warn( + "task_metadata for a merged ProgramSetQuantumTaskResult is synthesized " + "from multiple underlying tasks; it does not reflect any one underlying task. " + "Use the original per-task results for true task metadata.", + stacklevel=2, + ) + return self._task_metadata + + @property + def num_executables(self) -> int: + """int: The total number of executables in this program set task.""" + return self._num_executables + + @property + def program_set(self) -> ProgramSet | None: + """ProgramSet | None: The program set that was run, if provided to the constructor.""" + return self._program_set @staticmethod def from_object( @@ -425,15 +508,14 @@ def merge( programs = [_binding_to_program(binding) for binding in program_set.entries] executable_indices = list(program_set.enumerate_executables()) - binding_executable_counts = [_count_executables(b) for b in program_set.entries] - shots_per_executable = results[0].entries[0].shots_per_executable + shots_per_executable = program_set.shots_per_executable buffer = [None] * total_executables for k, result in enumerate(results): _buffer_result( k=k, result=result, - map_k=index_map[k], + parent_indices=index_map[k], program_set=program_set, programs=programs, executable_indices=executable_indices, @@ -443,23 +525,23 @@ def merge( entries = [] start = 0 for binding_idx, binding in enumerate(program_set.entries): - count = binding_executable_counts[binding_idx] + count = _count_executables(binding) program = programs[binding_idx] observables = binding.observables if isinstance(binding, CircuitBinding) else None - entries.append( - CompositeEntry( - entries=buffer[start : start + count], - program=program, - inputs=CompositeEntry._get_inputs(program, observables), - observables=observables, - shots_per_executable=shots_per_executable, - additional_metadata=None, - ) + entry = CompositeEntry( + entries=buffer[start : start + count], + program=program, + inputs=CompositeEntry._get_inputs(program, observables), + observables=observables, + shots_per_executable=shots_per_executable, + additional_metadata=None, ) + entry._was_merged = True + entries.append(entry) start += count - metas = [r.task_metadata for r in results] - return ProgramSetQuantumTaskResult( + metas = [r._task_metadata for r in results] + merged = ProgramSetQuantumTaskResult( entries=entries, task_metadata=ProgramSetTaskMetadata( id=";".join(meta.id for meta in metas), # Better way to do this? @@ -484,6 +566,8 @@ def merge( num_executables=total_executables, program_set=program_set, ) + merged._was_merged = True + return merged def __len__(self): return len(self.entries) @@ -613,7 +697,7 @@ def _count_executables(binding: CircuitBinding | Circuit) -> int: def _buffer_result( k: int, result: ProgramSetQuantumTaskResult, - map_k: list[int], + parent_indices: list[int], program_set: ProgramSet, programs: list[Program], executable_indices: list[tuple[int, int, int]], @@ -622,12 +706,12 @@ def _buffer_result( j = 0 for composite in result.entries: for entry in composite.entries: - if j >= len(map_k): + if j >= len(parent_indices): raise ValueError( f"t=Task {result.task_metadata.id} at index {k} " "produced more executables than index map expects" ) - orig_idx = map_k[j] + orig_idx = parent_indices[j] binding_idx, ps_idx, obs_idx = executable_indices[orig_idx] buffer[orig_idx] = _convert_measured_entry( entry, @@ -637,10 +721,10 @@ def _buffer_result( obs_idx, ) j += 1 - if j != len(map_k): + if j != len(parent_indices): raise ValueError( f"Task {result.task_metadata.id} at index {k} produced {j} executables " - f"but index map expected {len(map_k)}" + f"but index map expected {len(parent_indices)}" ) diff --git a/test/unit_tests/braket/tasks/test_program_set_quantum_task_result.py b/test/unit_tests/braket/tasks/test_program_set_quantum_task_result.py index e3170daaf..c18b535b1 100644 --- a/test/unit_tests/braket/tasks/test_program_set_quantum_task_result.py +++ b/test/unit_tests/braket/tasks/test_program_set_quantum_task_result.py @@ -12,6 +12,7 @@ # language governing permissions and limitations under the License. import json +import warnings from unittest.mock import Mock, patch import numpy as np @@ -269,6 +270,7 @@ def test_local(result_local, metadata, execution_failure): with pytest.raises(ValueError): entry.expectation(0) assert entry.additional_metadata.simulatorMetadata.executionDuration == 50 + assert entry.shots_per_executable > 0 success = entry[0] assert isinstance(success, MeasuredEntry) @@ -548,7 +550,7 @@ def _build_sub_quantum_result(sub_program_set, programs_execs, shots_per_executa return ProgramSetQuantumTaskResult.from_object(wire, sub_program_set) -def test_from_multiple_single_sub_task_no_split_roundtrips(circuit_rx_parametrized_fixture): +def test_from_multiple_single_task_no_split_roundtrips(circuit_rx_parametrized_fixture): """If split returns [self], from_multiple should reproduce from_object's output.""" binding = CircuitBinding( circuit_rx_parametrized_fixture, @@ -556,7 +558,7 @@ def test_from_multiple_single_sub_task_no_split_roundtrips(circuit_rx_parametriz observables=10 * Z(0) + X(0) - 0.01 * Y(0) @ X(1), ) ps = ProgramSet(binding) - subs, index_map = ps.split(100) # fits, so one sub-task identical to ps. + subs, index_map = ps.split(100) # fits, so one task identical to ps. assert subs == [ps] # Build a ProgramSetQuantumTaskResult that represents running this ps: the wire @@ -618,8 +620,9 @@ def test_from_multiple_split_list_observables(circuit_rx_parametrized_fixture): assert isinstance(measured, MeasuredEntry) assert measured.observable == binding.observables[i] assert composite.inputs == ParameterSets({"theta": [0.12]}) - # task metadata was aggregated across sub-tasks. + # task metadata was aggregated across tasks. assert merged.num_executables == 4 + assert merged.program_set is ps assert merged.task_metadata.requestedShots == sum( r.task_metadata.requestedShots for r in sub_results ) @@ -631,7 +634,7 @@ def test_from_multiple_split_list_observables(circuit_rx_parametrized_fixture): def test_from_multiple_split_sum_hamiltonian_reconstructs_expectation( circuit_rx_parametrized_fixture, ): - """Splitting a Sum Hamiltonian across multiple sub-tasks and then merging must + """Splitting a Sum Hamiltonian across multiple tasks and then merging must reconstruct the full expectation value, because scatter+regroup feeds the original Sum back into ``_compute_expectations``.""" # Same fixture as existing test_observables_no_inputs (with known expectation). @@ -661,8 +664,8 @@ def test_from_multiple_split_sum_hamiltonian_reconstructs_expectation( def test_from_multiple_mixed_bindings_and_failures(circuit_rx_parametrized_fixture): - """A program set with multiple bindings, split across sub-tasks, with one - executable failing in a sub-task. Failures must land at the correct original + """A program set with multiple bindings, split across tasks, with one + executable failing in a task. Failures must land at the correct original position in the merged result.""" c1 = circuit_rx_parametrized_fixture c2 = Circuit().rx(0, FreeParameter("phi")) @@ -741,10 +744,10 @@ def test_from_multiple_validates_index_map_size(circuit_rx_parametrized_fixture) binding = CircuitBinding(circuit_rx_parametrized_fixture, input_sets={"theta": [0.1, 0.2]}) ps = ProgramSet(binding) sub_result = _build_sub_quantum_result(ps, [[_make_exec_result(0), _make_exec_result(1)]]) - # index_map has 1 entry for 1 sub-task, but size doesn't match ps.total_executables. + # index_map has 1 entry for 1 task, but size doesn't match ps.total_executables. with pytest.raises(ValueError, match="Index map covers 1"): ProgramSetQuantumTaskResult.merge([sub_result], ps, [[0]]) - # Sub-task count doesn't match index_map's length. + # Task count doesn't match index_map's length. with pytest.raises(ValueError, match="1 task results but 2 entries in index_map"): ProgramSetQuantumTaskResult.merge([sub_result], ps, [[0], [1]]) @@ -773,11 +776,11 @@ def test_from_multiple_with_plain_circuit_entries(): assert merged[0].entries[0].inputs is None -def test_from_multiple_rejects_sub_task_over_index_map(circuit_rx_parametrized_fixture): - """Sub-task has more executables than index_map[k] covers.""" +def test_from_multiple_rejects_task_over_index_map(circuit_rx_parametrized_fixture): + """Task has more executables than index_map[k] covers.""" binding = CircuitBinding(circuit_rx_parametrized_fixture, input_sets={"theta": [0.1, 0.2]}) ps = ProgramSet(binding) - # Sub-task reports 2 executables, but index_map says there's only 1. + # Task reports 2 executables, but index_map says there's only 1. sub_result = _build_sub_quantum_result(ps, [[_make_exec_result(0), _make_exec_result(1)]]) with pytest.raises(ValueError, match="produced more executables than index map"): ProgramSetQuantumTaskResult.merge( @@ -787,16 +790,110 @@ def test_from_multiple_rejects_sub_task_over_index_map(circuit_rx_parametrized_f ) -def test_from_multiple_rejects_sub_task_under_index_map(circuit_rx_parametrized_fixture): - """Sub-task has fewer executables than index_map[k] covers.""" +def test_from_multiple_rejects_task_under_index_map(circuit_rx_parametrized_fixture): + """Task has fewer executables than index_map[k] covers.""" binding = CircuitBinding(circuit_rx_parametrized_fixture, input_sets={"theta": [0.1, 0.2]}) ps = ProgramSet(binding) - # Sub-task reports only 1 executable, but index_map says there are 2. + # Task reports only 1 executable, but index_map says there are 2. sub_result = _build_sub_quantum_result(ps, [[_make_exec_result(0)]]) with pytest.raises(ValueError, match="expected 2"): ProgramSetQuantumTaskResult.merge([sub_result], ps, [[0, 1]]) +def test_from_object_task_metadata_does_not_warn(circuit_rx_parametrized_fixture): + """A result built from a single ProgramSetTaskResult is not merged, so + accessing task_metadata should not emit a warning.""" + binding = CircuitBinding(circuit_rx_parametrized_fixture, input_sets={"theta": [0.1, 0.2]}) + ps = ProgramSet(binding) + result = _build_sub_quantum_result(ps, [[_make_exec_result(0), _make_exec_result(1)]]) + with warnings.catch_warnings(): + warnings.simplefilter("error") # turn any UserWarning into an exception + # Accessing task_metadata multiple times must not raise. + _ = result.task_metadata + _ = result.task_metadata.requestedShots + + +def test_merge_task_metadata_warns_on_every_access(circuit_rx_parametrized_fixture): + """Every access to task_metadata on a merged result emits a warning.""" + binding = CircuitBinding(circuit_rx_parametrized_fixture, input_sets={"theta": [0.1, 0.2]}) + ps = ProgramSet(binding) + subs, index_map = ps.split(1) + sub_results = [_build_sub_quantum_result(sub, [[_make_exec_result(0)]]) for sub in subs] + merged = ProgramSetQuantumTaskResult.merge(sub_results, ps, index_map) + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + _ = merged.task_metadata + _ = merged.task_metadata + _ = merged.task_metadata.id + user_warnings = [w for w in caught if issubclass(w.category, UserWarning)] + assert len(user_warnings) == 3 + for w in user_warnings: + assert "synthesized from multiple underlying tasks" in str(w.message) + + +def test_from_object_composite_additional_metadata_does_not_warn(circuit_rx_parametrized_fixture): + """A CompositeEntry on a non-merged result returns its actual additional_metadata + without warning.""" + binding = CircuitBinding(circuit_rx_parametrized_fixture, input_sets={"theta": [0.1, 0.2]}) + ps = ProgramSet(binding) + result = _build_sub_quantum_result(ps, [[_make_exec_result(0), _make_exec_result(1)]]) + composite = result[0] + with warnings.catch_warnings(): + warnings.simplefilter("error") # any UserWarning becomes a test failure + meta = composite.additional_metadata + _ = composite.additional_metadata # repeat read also silent + assert meta is not None # populated from the wire fixture + + +def test_merge_composite_additional_metadata_warns_on_every_access( + circuit_rx_parametrized_fixture, +): + """Every access to additional_metadata on a merged CompositeEntry emits a + warning, and the value is None.""" + binding = CircuitBinding(circuit_rx_parametrized_fixture, input_sets={"theta": [0.1, 0.2]}) + ps = ProgramSet(binding) + subs, index_map = ps.split(1) + sub_results = [_build_sub_quantum_result(sub, [[_make_exec_result(0)]]) for sub in subs] + merged = ProgramSetQuantumTaskResult.merge(sub_results, ps, index_map) + composite = merged[0] + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + a = composite.additional_metadata + b = composite.additional_metadata + user_warnings = [w for w in caught if issubclass(w.category, UserWarning)] + assert len(user_warnings) == 2 + for w in user_warnings: + assert "additional_metadata for a CompositeEntry on a merged" in str(w.message) + assert a is None and b is None + + +def test_merge_does_not_warn_when_reading_inputs_during_construction( + circuit_rx_parametrized_fixture, +): + """If an already-merged result is fed back into ``merge`` as an input, the + internal read of ``task_metadata`` to build the aggregated metadata must not + trigger that input's task_metadata warning.""" + binding = CircuitBinding(circuit_rx_parametrized_fixture, input_sets={"theta": [0.1, 0.2]}) + ps = ProgramSet(binding) + subs, index_map = ps.split(1) + sub_results = [_build_sub_quantum_result(sub, [[_make_exec_result(0)]]) for sub in subs] + inner_merged = ProgramSetQuantumTaskResult.merge(sub_results, ps, index_map) + + # Re-merging a single already-merged result is degenerate but exercises the + # internal ``task_metadata`` reads inside merge. + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + ProgramSetQuantumTaskResult.merge([inner_merged], ps, [list(range(ps.total_executables))]) + user_warnings = [w for w in caught if issubclass(w.category, UserWarning)] + # No warning should fire from inside merge itself. + assert user_warnings == [], ( + f"merge emitted unexpected warnings during construction: " + f"{[str(w.message) for w in user_warnings]}" + ) + + def ghz_test(n): """Local ghz helper so tests don't depend on program_set_test_utils.""" circuit = Circuit().h(0) From 098ee062b8dcfdfaa2ba950f5ec0461455521a49 Mon Sep 17 00:00:00 2001 From: Cody Wang Date: Mon, 15 Jun 2026 21:57:40 -0700 Subject: [PATCH 3/3] Update program_set_quantum_task_result.py --- .../tasks/program_set_quantum_task_result.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/braket/tasks/program_set_quantum_task_result.py b/src/braket/tasks/program_set_quantum_task_result.py index 38e8b15c9..71b611cb0 100644 --- a/src/braket/tasks/program_set_quantum_task_result.py +++ b/src/braket/tasks/program_set_quantum_task_result.py @@ -510,16 +510,16 @@ def merge( executable_indices = list(program_set.enumerate_executables()) shots_per_executable = program_set.shots_per_executable - buffer = [None] * total_executables + executable_results_in_order = [None] * total_executables for k, result in enumerate(results): - _buffer_result( + _reorder_executable_results( k=k, result=result, parent_indices=index_map[k], program_set=program_set, programs=programs, executable_indices=executable_indices, - buffer=buffer, + out=executable_results_in_order, ) entries = [] @@ -529,7 +529,7 @@ def merge( program = programs[binding_idx] observables = binding.observables if isinstance(binding, CircuitBinding) else None entry = CompositeEntry( - entries=buffer[start : start + count], + entries=executable_results_in_order[start : start + count], program=program, inputs=CompositeEntry._get_inputs(program, observables), observables=observables, @@ -694,14 +694,14 @@ def _count_executables(binding: CircuitBinding | Circuit) -> int: return num_ps * num_obs -def _buffer_result( +def _reorder_executable_results( k: int, result: ProgramSetQuantumTaskResult, parent_indices: list[int], program_set: ProgramSet, programs: list[Program], executable_indices: list[tuple[int, int, int]], - buffer: list[MeasuredEntry | ProgramSetExecutableFailure | None], + out: list[MeasuredEntry | ProgramSetExecutableFailure | None], ) -> None: j = 0 for composite in result.entries: @@ -713,7 +713,7 @@ def _buffer_result( ) orig_idx = parent_indices[j] binding_idx, ps_idx, obs_idx = executable_indices[orig_idx] - buffer[orig_idx] = _convert_measured_entry( + out[orig_idx] = _convert_measured_entry( entry, program_set.entries[binding_idx], programs[binding_idx],