fix: Force StateVector and DensityMatrix values to be ndarrays and test

src/braket/simulator_v2/simulator.py

@@ -1,5 +1,6 @@
 import sys
 
+import numpy as np
 from braket.default_simulator.operation_helpers import from_braket_instruction
 from braket.default_simulator.result_types import TargetedResultType
 from braket.default_simulator.simulator import BaseLocalSimulator
@@ -8,13 +9,32 @@
     GateModelSimulatorDeviceCapabilities,
     GateModelSimulatorDeviceParameters,
 )
+from braket.ir.jaqcd import DensityMatrix
 from braket.ir.jaqcd import Program as JaqcdProgram
+from braket.ir.jaqcd import StateVector
 from braket.ir.openqasm import Program as OpenQASMProgram
 from braket.task_result import GateModelTaskResult
 
 from .julia_import import jl, jlBraketSimulator
 
 
+def _analytic_result_value_to_ndarray(
+    task_result: GateModelTaskResult,
+) -> GateModelTaskResult:
+    """Convert any StateVector or DensityMatrix result values from raw Python lists to the expected
+    np.ndarray. This must be done because the wrapper Julia simulator results Python lists to comply
+    with the pydantic specification for ResultTypeValues.
+    """
+    for result_ind, result_type in enumerate(task_result.resultTypes):
+        if isinstance(result_type.type, StateVector) or isinstance(
+            result_type.type, DensityMatrix
+        ):
+            task_result.resultTypes[result_ind].value = np.asarray(
+                task_result.resultTypes[result_ind].value
+            )
+    return task_result
+
+
 class StateVectorSimulatorV2(BaseLocalSimulator):
     """A state vector simulator meant to run directly on the user's machine using a Julia backend.
 
@@ -88,6 +108,8 @@ def run_jaqcd(
             )
         r = jl.simulate(self._device, [circuit_ir], qubit_count, shots)
         r.additionalMetadata.action = circuit_ir
+        if not shots:
+            r = _analytic_result_value_to_ndarray(r)
         return r
 
     def run_openqasm(
@@ -153,9 +175,11 @@ def run_openqasm(
             self._device, [circuit], qubit_count, shots, measured_qubits=measured_qubits
         )
         r.additionalMetadata.action = openqasm_ir
-        # attach the result types
         if shots:
+            # attach the result types
             r.resultTypes = results
+        else:
+            r = _analytic_result_value_to_ndarray(r)
         return r
 
     @property
@@ -473,6 +497,8 @@ def run_jaqcd(
             )
         r = jl.simulate(self._device, [circuit_ir], qubit_count, shots)
         r.additionalMetadata.action = circuit_ir
+        if not shots:
+            r = _analytic_result_value_to_ndarray(r)
         return r
 
     def run_openqasm(
@@ -540,6 +566,8 @@ def run_openqasm(
         # attach the result types
         if shots:
             r.resultTypes = results
+        else:
+            r = _analytic_result_value_to_ndarray(r)
         return r
 
     @property

test/unit_tests/braket/simulator_v2/test_density_matrix_simulator.py

@@ -22,7 +22,7 @@
     GateModelSimulatorDeviceCapabilities,
     GateModelSimulatorDeviceParameters,
 )
-from braket.ir.jaqcd import Expectation
+from braket.ir.jaqcd import DensityMatrix, Expectation
 from braket.ir.jaqcd import Program as JaqcdProgram
 from braket.ir.openqasm import Program as OpenQASMProgram
 from braket.task_result import AdditionalMetadata, TaskMetadata
@@ -846,3 +846,34 @@ def test_measure_targets():
     assert 400 < np.sum(measurements, axis=0)[0] < 600
     assert len(measurements[0]) == 1
     assert result.measuredQubits == [0]
+
+
+@pytest.mark.parametrize(
+    "jaqcd_string, oq3_pragma, jaqcd_type",
+    [
+        ["densitymatrix", "density_matrix", DensityMatrix()],
+    ],
+)
+def test_simulator_analytic_value_type(jaqcd_string, oq3_pragma, jaqcd_type):
+    simulator = DensityMatrixSimulator()
+    jaqcd = JaqcdProgram.parse_raw(
+        json.dumps(
+            {
+                "instructions": [{"type": "h", "target": 0}],
+                "results": [{"type": jaqcd_string}],
+            }
+        )
+    )
+    qasm = OpenQASMProgram(
+        source=f"""
+        qubit q;
+        h q;
+        #pragma braket result {oq3_pragma}
+        """
+    )
+    result = simulator.run(jaqcd, qubit_count=2, shots=0)
+    assert result.resultTypes[0].type == jaqcd_type
+    assert isinstance(result.resultTypes[0].value, np.ndarray)
+    result = simulator.run(qasm, shots=0)
+    assert result.resultTypes[0].type == jaqcd_type
+    assert isinstance(result.resultTypes[0].value, np.ndarray)

test/unit_tests/braket/simulator_v2/test_state_vector_simulator.py

@@ -1410,3 +1410,35 @@ def test_rotation_parameter_expressions(operation, state_vector):
     result = simulator.run(OpenQASMProgram(source=qasm), shots=0)
     assert result.resultTypes[0].type == StateVector()
     assert np.allclose(result.resultTypes[0].value, np.array(state_vector))
+
+
+@pytest.mark.parametrize(
+    "jaqcd_string, oq3_pragma, jaqcd_type",
+    [
+        ["statevector", "state_vector", StateVector()],
+        ["densitymatrix", "density_matrix", DensityMatrix()],
+    ],
+)
+def test_simulator_analytic_value_type(jaqcd_string, oq3_pragma, jaqcd_type):
+    simulator = StateVectorSimulator()
+    jaqcd = JaqcdProgram.parse_raw(
+        json.dumps(
+            {
+                "instructions": [{"type": "h", "target": 0}],
+                "results": [{"type": jaqcd_string}],
+            }
+        )
+    )
+    qasm = OpenQASMProgram(
+        source=f"""
+        qubit q;
+        h q;
+        #pragma braket result {oq3_pragma}
+        """
+    )
+    result = simulator.run(jaqcd, qubit_count=2, shots=0)
+    assert result.resultTypes[0].type == jaqcd_type
+    assert isinstance(result.resultTypes[0].value, np.ndarray)
+    result = simulator.run(qasm, shots=0)
+    assert result.resultTypes[0].type == jaqcd_type
+    assert isinstance(result.resultTypes[0].value, np.ndarray)