fix: have hamiltonian updates inline

src/braket/analog_hamiltonian_simulator/rydberg/rydberg_simulator_helpers.py

@@ -558,11 +558,9 @@ def _get_hamiltonian(
     for rabi_op, rabi_coef, detuning_op, detuning_coef in zip(
         rabi_ops, rabi_coefs, detuning_ops, detuning_coefs
     ):
-        hamiltonian += (
-            rabi_op * rabi_coef[index_time] / 2
-            + (rabi_op.T.conj() * np.conj(rabi_coef[index_time]) / 2)
-            - detuning_op * detuning_coef[index_time]
-        )
+        hamiltonian += rabi_op * rabi_coef[index_time] / 2
+        hamiltonian += rabi_op.T.conj() * np.conj(rabi_coef[index_time]) / 2
+        hamiltonian -= detuning_op * detuning_coef[index_time]
 
     # Add local detuning
     for local_detuning_op, local_detuning_coef in zip(local_detuning_ops, local_detuing_coefs):

src/braket/default_simulator/linalg_utils.py

@@ -46,14 +46,15 @@ def multiply_matrix(
     Returns:
         np.ndarray: The state after the matrix has been applied.
     """
+    targets = np.array(targets)
     if not controls:
         return _multiply_matrix(state, matrix, targets)
-
     control_state = control_state or (1,) * len(controls)
+
     num_qubits = len(state.shape)
     control_slices = {i: _SLICES[state] for i, state in zip(controls, control_state)}
     ctrl_index = tuple(
-        control_slices[i] if i in controls else _NO_CONTROL_SLICE for i in range(num_qubits)
+        control_slices[i] if i in control_slices else _NO_CONTROL_SLICE for i in range(num_qubits)
     )
     state[ctrl_index] = _multiply_matrix(state[ctrl_index], matrix, targets)
     return state
@@ -62,7 +63,7 @@ def multiply_matrix(
 def _multiply_matrix(
     state: np.ndarray,
     matrix: np.ndarray,
-    targets: tuple[int, ...],
+    targets: np.ndarray,
 ) -> np.ndarray:
     """Multiplies the given matrix by the given state, applying the matrix on the target qubits.
 
@@ -74,13 +75,14 @@ def _multiply_matrix(
     Returns:
         np.ndarray: The state after the matrix has been applied.
     """
-    gate_matrix = np.reshape(matrix, [2] * len(targets) * 2)
+    num_targets = len(targets)
+    gate_matrix = np.reshape(matrix, [2] * num_targets * 2)
     axes = (
-        np.arange(len(targets), 2 * len(targets)),
+        np.arange(num_targets, 2 * num_targets),
         targets,
     )
-    product = np.tensordot(gate_matrix, state, axes=axes)
 
+    product = np.tensordot(gate_matrix, state, axes=axes)
     # Axes given in `operation.targets` are in the first positions.
     unused_idxs = [idx for idx in range(len(state.shape)) if idx not in targets]
     permutation = list(targets) + unused_idxs

src/braket/default_simulator/simulator.py

@@ -473,9 +473,10 @@ def run_openqasm(
         simulation = self.initialize_simulation(
             qubit_count=qubit_count, shots=shots, batch_size=batch_size
         )
-        simulation.evolve(operations)
-
-        if not shots:
+        if shots:
+            simulation.evolve(operations + circuit.basis_rotation_instructions)
+        else:
+            simulation.evolve(operations)
             result_types = BaseLocalSimulator._translate_result_types(results)
             BaseLocalSimulator._validate_result_types_qubits_exist(
                 [
@@ -490,8 +491,6 @@ def run_openqasm(
                 result_types,
                 simulation,
             )
-        else:
-            simulation.evolve(circuit.basis_rotation_instructions)
 
         return self._create_results_obj(results, openqasm_ir, simulation, measured_qubits)