Moving state_vector method to task (#253)

+ Adding a bug fix for `pauli_expectation`.

---------

Co-authored-by: Copilot <[email protected]>

Author: weinbe58

src/bloqade/pyqrack/device.py

@@ -60,9 +60,7 @@ def state_vector(
         kwargs: dict[str, Any] | None = None,
     ) -> list[complex]:
         """Runs task and returns the state vector."""
-        task = self.task(kernel, args, kwargs)
-        task.run()
-        return task.state.sim_reg.out_ket()
+        return self.task(kernel, args, kwargs).state_vector()
 
     @staticmethod
     def pauli_expectation(pauli: list[Pauli], qubits: list[PyQrackQubit]) -> float:
@@ -96,7 +94,7 @@ def pauli_expectation(pauli: list[Pauli], qubits: list[PyQrackQubit]) -> float:
         if len(qubit_ids) != len(set(qubit_ids)):
             raise ValueError("Qubits must be unique.")
 
-        return sim_reg.pauli_expectation(pauli, qubit_ids)
+        return sim_reg.pauli_expectation(qubit_ids, pauli)
 
 
 @dataclass

src/bloqade/pyqrack/task.py

@@ -1,4 +1,4 @@
-from typing import TypeVar, ParamSpec
+from typing import TypeVar, ParamSpec, cast
 from dataclasses import dataclass
 
 from bloqade.task import AbstractSimulatorTask
@@ -19,12 +19,20 @@ class PyQrackSimulatorTask(AbstractSimulatorTask[Param, RetType, MemoryType]):
     pyqrack_interp: PyQrackInterpreter[MemoryType]
 
     def run(self) -> RetType:
-        return self.pyqrack_interp.run(
-            self.kernel,
-            args=self.args,
-            kwargs=self.kwargs,
+        return cast(
+            RetType,
+            self.pyqrack_interp.run(
+                self.kernel,
+                args=self.args,
+                kwargs=self.kwargs,
+            ),
         )
 
     @property
     def state(self) -> MemoryType:
         return self.pyqrack_interp.memory
+
+    def state_vector(self) -> list[complex]:
+        """Returns the state vector of the simulator."""
+        self.run()
+        return self.state.sim_reg.out_ket()