From 58e898a7e5604a3e03be04a9d5fbf771a5f6d7e2 Mon Sep 17 00:00:00 2001 From: Cody Wang Date: Sun, 10 May 2026 14:20:57 -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 9f7084c97..489b68dc5 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..886a4a22e 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, mapping = 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, mapping) + + 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, mapping = 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, mapping) + 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, mapping = 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, mapping) + 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, mapping = 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 mapping[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, mapping) + 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_mapping_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)]]) + # mapping 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 mapping'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, mapping = 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, mapping) + 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_mapping(circuit_rx_parametrized_fixture): + """Sub-task has more executables than mapping[k] covers.""" + binding = CircuitBinding(circuit_rx_parametrized_fixture, input_sets={"theta": [0.1, 0.2]}) + ps = ProgramSet(binding) + # Sub-task reports 2 executables, but mapping 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_mapping(circuit_rx_parametrized_fixture): + """Sub-task has fewer executables than mapping[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 mapping 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 cb558790f89507d5a32bce0aeb90b102fa732401 Mon Sep 17 00:00:00 2001 From: Cody Wang Date: Sun, 10 May 2026 14:27:38 -0700 Subject: [PATCH 2/3] reformat --- src/braket/tasks/program_set_quantum_task_result.py | 4 ++-- 1 file changed, 2 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 9f7084c97..c4c37b00e 100644 --- a/src/braket/tasks/program_set_quantum_task_result.py +++ b/src/braket/tasks/program_set_quantum_task_result.py @@ -260,7 +260,7 @@ def _get_inputs(program: Program, observables: Sum | list[Observable] | None) -> def _get_executable_results( executable_results: Sequence[ ProgramSetExecutableResult | ProgramSetExecutableFailure | str - ], + ], program: Program, observables: Sum | list[Observable] | None, shots_per_executable: int, @@ -300,7 +300,7 @@ def _dispatch_executable_result( program=program.source, shots=shots_per_executable, inputs={k: v[result.inputsIndex] for k, v in (program.inputs or {}).items()} - or None, + or None, observable=( observables[result.inputsIndex % len(observables)] if observables else None ), From 94187237cf21240d6b6caef8c3142b771f29b2dd Mon Sep 17 00:00:00 2001 From: Cody Wang Date: Sun, 10 May 2026 14:20:57 -0700 Subject: [PATCH 3/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..886a4a22e 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, mapping = 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, mapping) + + 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, mapping = 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, mapping) + 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, mapping = 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, mapping) + 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, mapping = 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 mapping[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, mapping) + 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_mapping_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)]]) + # mapping 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 mapping'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, mapping = 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, mapping) + 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_mapping(circuit_rx_parametrized_fixture): + """Sub-task has more executables than mapping[k] covers.""" + binding = CircuitBinding(circuit_rx_parametrized_fixture, input_sets={"theta": [0.1, 0.2]}) + ps = ProgramSet(binding) + # Sub-task reports 2 executables, but mapping 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_mapping(circuit_rx_parametrized_fixture): + """Sub-task has fewer executables than mapping[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 mapping 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