bloqade-circuit 0.7.4__py3-none-any.whl → 0.7.6__py3-none-any.whl

bloqade/native/__init__.py

@@ -0,0 +1,24 @@
+from .stdlib import broadcast as broadcast
+from ._prelude import kernel as kernel
+from .stdlib.simple import (
+    h as h,
+    s as s,
+    t as t,
+    x as x,
+    y as y,
+    z as z,
+    cx as cx,
+    cy as cy,
+    cz as cz,
+    rx as rx,
+    ry as ry,
+    rz as rz,
+    u3 as u3,
+    rot as rot,
+    s_dag as s_dag,
+    shift as shift,
+    sqrt_x as sqrt_x,
+    sqrt_y as sqrt_y,
+    sqrt_x_adj as sqrt_x_adj,
+    sqrt_y_adj as sqrt_y_adj,
+)

bloqade/native/_prelude.py

@@ -0,0 +1,45 @@
+from typing import Annotated
+
+from kirin import ir
+from kirin.passes import Default
+from kirin.prelude import structural_no_opt
+from typing_extensions import Doc
+
+from bloqade.squin import qubit
+
+from .dialects import gates
+
+
+@ir.dialect_group(structural_no_opt.union([gates, qubit]))
+def kernel(self):
+    """Compile a function to a native kernel."""
+
+    def run_pass(
+        mt,
+        *,
+        verify: Annotated[
+            bool, Doc("run `verify` before running passes, default is `True`")
+        ] = True,
+        typeinfer: Annotated[
+            bool,
+            Doc(
+                "run type inference and apply the inferred type to IR, default `False`"
+            ),
+        ] = False,
+        fold: Annotated[bool, Doc("run folding passes")] = True,
+        aggressive: Annotated[
+            bool, Doc("run aggressive folding passes if `fold=True`")
+        ] = False,
+        no_raise: Annotated[bool, Doc("do not raise exception during analysis")] = True,
+    ) -> None:
+        default_pass = Default(
+            self,
+            verify=verify,
+            fold=fold,
+            aggressive=aggressive,
+            typeinfer=typeinfer,
+            no_raise=no_raise,
+        )
+        default_pass.fixpoint(mt)
+
+    return run_pass

bloqade/native/dialects/gates/__init__.py

@@ -0,0 +1,3 @@
+from .stmts import CZ as CZ, R as R, Rz as Rz
+from ._dialect import dialect as dialect
+from ._interface import r as r, cz as cz, rz as rz

bloqade/native/dialects/gates/_dialect.py

@@ -0,0 +1,3 @@
+from kirin import ir
+
+dialect = ir.Dialect("bloqade.native")

bloqade/native/dialects/gates/_interface.py

@@ -0,0 +1,32 @@
+import typing
+
+from kirin import lowering
+from kirin.dialects import ilist
+
+from bloqade.squin import qubit
+
+from .stmts import CZ, R, Rz
+
+Len = typing.TypeVar("Len")
+
+
+@lowering.wraps(CZ)
+def cz(
+    ctrls: ilist.IList[qubit.Qubit, Len],
+    qargs: ilist.IList[qubit.Qubit, Len],
+): ...
+
+
+@lowering.wraps(R)
+def r(
+    inputs: ilist.IList[qubit.Qubit, typing.Any],
+    axis_angle: float,
+    rotation_angle: float,
+): ...
+
+
+@lowering.wraps(Rz)
+def rz(
+    inputs: ilist.IList[qubit.Qubit, typing.Any],
+    rotation_angle: float,
+): ...

bloqade/native/dialects/gates/stmts.py

@@ -0,0 +1,31 @@
+from kirin import ir, types, lowering
+from kirin.decl import info, statement
+from kirin.dialects import ilist
+
+from bloqade.squin import qubit
+
+from ._dialect import dialect
+
+N = types.TypeVar("N")
+
+
+@statement(dialect=dialect)
+class CZ(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    ctrls: ir.SSAValue = info.argument(ilist.IListType[qubit.QubitType, N])
+    qargs: ir.SSAValue = info.argument(ilist.IListType[qubit.QubitType, N])
+
+
+@statement(dialect=dialect)
+class R(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    inputs: ir.SSAValue = info.argument(ilist.IListType[qubit.QubitType, types.Any])
+    axis_angle: ir.SSAValue = info.argument(types.Float)
+    rotation_angle: ir.SSAValue = info.argument(types.Float)
+
+
+@statement(dialect=dialect)
+class Rz(ir.Statement):
+    traits = frozenset({lowering.FromPythonCall()})
+    inputs: ir.SSAValue = info.argument(ilist.IListType[qubit.QubitType, types.Any])
+    rotation_angle: ir.SSAValue = info.argument(types.Float)

bloqade/native/stdlib/broadcast.py

@@ -0,0 +1,246 @@
+import math
+from typing import Any, TypeVar
+
+from kirin.dialects import ilist
+
+from bloqade.squin import qubit
+from bloqade.native._prelude import kernel
+from bloqade.native.dialects.gates import _interface as native
+
+
+@kernel
+def _radian_to_turn(angle: float) -> float:
+    """Convert an angle from radians to turns.
+
+    Args:
+        angle (float): Angle in radians.
+
+    Returns:
+        float: Equivalent angle in turns.
+    """
+    return angle / (2 * math.pi)
+
+
+@kernel
+def rx(angle: float, qubits: ilist.IList[qubit.Qubit, Any]):
+    """Apply an RX rotation gate on a group of qubits.
+
+    Args:
+        angle (float): Rotation angle in radians.
+        qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
+    """
+    native.r(qubits, 0.0, _radian_to_turn(angle))
+
+
+@kernel
+def x(qubits: ilist.IList[qubit.Qubit, Any]):
+    """Apply a Pauli-X gate on a group of qubits.
+
+    Args:
+        qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
+    """
+    rx(math.pi, qubits)
+
+
+@kernel
+def sqrt_x(qubits: ilist.IList[qubit.Qubit, Any]):
+    """Apply a sqrt(X) gate on a group of qubits.
+
+    Args:
+        qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
+    """
+    rx(math.pi / 2.0, qubits)
+
+
+@kernel
+def sqrt_x_adj(qubits: ilist.IList[qubit.Qubit, Any]):
+    """Apply the adjoint of sqrt(X) on a group of qubits.
+
+    Args:
+        qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
+    """
+    rx(-math.pi / 2.0, qubits)
+
+
+@kernel
+def ry(angle: float, qubits: ilist.IList[qubit.Qubit, Any]):
+    """Apply an RY rotation gate on a group of qubits.
+
+    Args:
+        angle (float): Rotation angle in radians.
+        qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
+    """
+    native.r(qubits, 0.25, _radian_to_turn(angle))
+
+
+@kernel
+def y(qubits: ilist.IList[qubit.Qubit, Any]):
+    """Apply a Pauli-Y gate on a group of qubits.
+
+    Args:
+        qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
+    """
+    ry(math.pi, qubits)
+
+
+@kernel
+def sqrt_y(qubits: ilist.IList[qubit.Qubit, Any]):
+    """Apply a sqrt(Y) gate on a group of qubits.
+
+    Args:
+        qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
+    """
+    ry(-math.pi / 2.0, qubits)
+
+
+@kernel
+def sqrt_y_adj(qubits: ilist.IList[qubit.Qubit, Any]):
+    """Apply the adjoint of sqrt(Y) on a group of qubits.
+
+    Args:
+        qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
+    """
+    ry(math.pi / 2.0, qubits)
+
+
+@kernel
+def rz(angle: float, qubits: ilist.IList[qubit.Qubit, Any]):
+    """Apply an RZ rotation gate on a group of qubits.
+
+    Args:
+        angle (float): Rotation angle in radians.
+        qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
+    """
+    native.rz(qubits, _radian_to_turn(angle))
+
+
+@kernel
+def z(qubits: ilist.IList[qubit.Qubit, Any]):
+    """Apply a Pauli-Z gate on a group of qubits.
+
+    Args:
+        qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
+    """
+    rz(math.pi, qubits)
+
+
+@kernel
+def s(qubits: ilist.IList[qubit.Qubit, Any]):
+    """Apply an S gate on a group of qubits.
+
+    Args:
+        qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
+    """
+    rz(math.pi / 2.0, qubits)
+
+
+@kernel
+def s_adj(qubits: ilist.IList[qubit.Qubit, Any]):
+    """Apply the adjoint of the S gate on a group of qubits.
+
+    Args:
+        qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
+    """
+    rz(-math.pi / 2.0, qubits)
+
+
+@kernel
+def h(qubits: ilist.IList[qubit.Qubit, Any]):
+    """Apply a Hadamard gate on a group of qubits.
+
+    Args:
+        qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
+    """
+    s(qubits)
+    sqrt_x(qubits)
+    s(qubits)
+
+
+@kernel
+def t(qubits: ilist.IList[qubit.Qubit, Any]):
+    """Apply a T gate on a group of qubits.
+
+    Args:
+        qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
+    """
+    rz(math.pi / 4.0, qubits)
+
+
+@kernel
+def shift(angle: float, qubits: ilist.IList[qubit.Qubit, Any]):
+    """Apply a phase shift to the |1> state on a group of qubits.
+
+    Args:
+        angle (float): Phase shift angle in radians.
+        qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
+    """
+    rz(angle / 2.0, qubits)
+
+
+@kernel
+def rot(phi: float, theta: float, omega: float, qubits: ilist.IList[qubit.Qubit, Any]):
+    """Apply a general single-qubit rotation on a group of qubits.
+
+    Args:
+        phi (float): Z rotation before Y (radians).
+        theta (float): Y rotation (radians).
+        omega (float): Z rotation after Y (radians).
+        qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
+    """
+    rz(phi, qubits)
+    ry(theta, qubits)
+    rz(omega, qubits)
+
+
+@kernel
+def u3(theta: float, phi: float, lam: float, qubits: ilist.IList[qubit.Qubit, Any]):
+    """Apply the U3 gate on a group of qubits.
+
+    Args:
+        theta (float): Rotation around Y axis (radians).
+        phi (float): Global phase shift component (radians).
+        lam (float): Z rotations in decomposition (radians).
+        qubits (ilist.IList[qubit.Qubit, Any]): Target qubits.
+    """
+    rot(lam, theta, -lam, qubits)
+    shift(phi + lam, qubits)
+
+
+N = TypeVar("N")
+
+
+@kernel
+def cz(controls: ilist.IList[qubit.Qubit, N], qubits: ilist.IList[qubit.Qubit, N]):
+    """Apply a controlled-Z gate on a pairs of qubits.
+
+    Args:
+        controls (ilist.IList[qubit.Qubit, N]): Control qubits.
+        qubits (ilist.IList[qubit.Qubit, N]): Target qubits.
+    """
+    native.cz(controls, qubits)
+
+
+@kernel
+def cx(controls: ilist.IList[qubit.Qubit, N], targets: ilist.IList[qubit.Qubit, N]):
+    """Apply a controlled-X gate on a pairs of qubits.
+
+    Args:
+        controls (ilist.IList[qubit.Qubit, N]): Control qubits.
+        targets (ilist.IList[qubit.Qubit, N]): Target qubits.
+    """
+    sqrt_y_adj(targets)
+    cz(controls, targets)
+    sqrt_y(targets)
+
+
+@kernel
+def cy(controls: ilist.IList[qubit.Qubit, N], targets: ilist.IList[qubit.Qubit, N]):
+    """Apply a controlled-Y gate on a pairs of qubits.
+
+    Args:
+        controls (ilist.IList[qubit.Qubit, N]): Control qubits.
+        targets (ilist.IList[qubit.Qubit, N]): Target qubits.
+    """
+    sqrt_x(targets)
+    cz(controls, targets)
+    sqrt_x_adj(targets)

bloqade/native/stdlib/simple.py

@@ -0,0 +1,219 @@
+from kirin.dialects import ilist
+
+from bloqade.squin import qubit
+
+from . import broadcast
+from .._prelude import kernel
+
+
+@kernel
+def rx(angle: float, qubit: qubit.Qubit):
+    """Apply an RX rotation gate on a single qubit.
+
+    Args:
+        angle (float): Rotation angle in radians.
+        qubit (qubit.Qubit): The qubit to apply the rotation to.
+    """
+    broadcast.rx(angle, ilist.IList([qubit]))
+
+
+@kernel
+def x(qubit: qubit.Qubit):
+    """Apply a Pauli-X gate on a single qubit.
+
+    Args:
+        qubit (qubit.Qubit): The qubit to apply the X gate to.
+    """
+    broadcast.x(ilist.IList([qubit]))
+
+
+@kernel
+def sqrt_x(qubit: qubit.Qubit):
+    """Apply a sqrt(X) gate on a single qubit.
+
+    Args:
+        qubit (qubit.Qubit): The qubit to apply the sqrt(X) gate to.
+    """
+    broadcast.sqrt_x(ilist.IList([qubit]))
+
+
+@kernel
+def sqrt_x_adj(qubit: qubit.Qubit):
+    """Apply the adjoint of sqrt(X) on a single qubit.
+
+    Args:
+        qubit (qubit.Qubit): The qubit to apply the adjoint sqrt(X) gate to.
+    """
+    broadcast.sqrt_x_adj(ilist.IList([qubit]))
+
+
+@kernel
+def ry(angle: float, qubit: qubit.Qubit):
+    """Apply an RY rotation gate on a single qubit.
+
+    Args:
+        angle (float): Rotation angle in radians.
+        qubit (qubit.Qubit): The qubit to apply the rotation to.
+    """
+    broadcast.ry(angle, ilist.IList([qubit]))
+
+
+@kernel
+def y(qubit: qubit.Qubit):
+    """Apply a Pauli-Y gate on a single qubit.
+
+    Args:
+        qubit (qubit.Qubit): The qubit to apply the Y gate to.
+    """
+    broadcast.y(ilist.IList([qubit]))
+
+
+@kernel
+def sqrt_y(qubit: qubit.Qubit):
+    """Apply a sqrt(Y) gate on a single qubit.
+
+    Args:
+        qubit (qubit.Qubit): The qubit to apply the sqrt(Y) gate to.
+    """
+    broadcast.sqrt_y(ilist.IList([qubit]))
+
+
+@kernel
+def sqrt_y_adj(qubit: qubit.Qubit):
+    """Apply the adjoint of sqrt(Y) on a single qubit.
+
+    Args:
+        qubit (qubit.Qubit): The qubit to apply the adjoint sqrt(Y) gate to.
+    """
+    broadcast.sqrt_y_adj(ilist.IList([qubit]))
+
+
+@kernel
+def rz(angle: float, qubit: qubit.Qubit):
+    """Apply an RZ rotation gate on a single qubit.
+
+    Args:
+        angle (float): Rotation angle in radians.
+        qubit (qubit.Qubit): The qubit to apply the rotation to.
+    """
+    broadcast.rz(angle, ilist.IList([qubit]))
+
+
+@kernel
+def z(qubit: qubit.Qubit):
+    """Apply a Pauli-Z gate on a single qubit.
+
+    Args:
+        qubit (qubit.Qubit): The qubit to apply the Z gate to.
+    """
+    broadcast.z(ilist.IList([qubit]))
+
+
+@kernel
+def s(qubit: qubit.Qubit):
+    """Apply an S gate on a single qubit.
+
+    Args:
+        qubit (qubit.Qubit): The qubit to apply the S gate to.
+    """
+    broadcast.s(ilist.IList([qubit]))
+
+
+@kernel
+def s_dag(qubit: qubit.Qubit):
+    """Apply the adjoint of the S gate on a single qubit.
+
+    Args:
+        qubit (qubit.Qubit): The qubit to apply the adjoint S gate to.
+    """
+    broadcast.s_adj(ilist.IList([qubit]))
+
+
+@kernel
+def h(qubit: qubit.Qubit):
+    """Apply a Hadamard gate on a single qubit.
+
+    Args:
+        qubit (qubit.Qubit): The qubit to apply the Hadamard gate to.
+    """
+    broadcast.h(ilist.IList([qubit]))
+
+
+@kernel
+def t(qubit: qubit.Qubit):
+    """Apply a T gate on a single qubit.
+
+    Args:
+        qubit (qubit.Qubit): The qubit to apply the T gate to.
+    """
+    broadcast.t(ilist.IList([qubit]))
+
+
+@kernel
+def shift(angle: float, qubit: qubit.Qubit):
+    """Apply a phase shift on the |1> state of a single qubit.
+
+    Args:
+        angle (float): Shift angle in radians.
+        qubit (qubit.Qubit): The qubit to apply the shift to.
+    """
+    broadcast.shift(angle, ilist.IList([qubit]))
+
+
+@kernel
+def rot(phi: float, theta: float, omega: float, qubit: qubit.Qubit):
+    """Apply a general single-qubit rotation on a single qubit.
+
+    Args:
+        phi (float): Z rotation before Y (radians).
+        theta (float): Y rotation (radians).
+        omega (float): Z rotation after Y (radians).
+        qubit (qubit.Qubit): The qubit to apply the rotation to.
+    """
+    broadcast.rot(phi, theta, omega, ilist.IList([qubit]))
+
+
+@kernel
+def u3(theta: float, phi: float, lam: float, qubit: qubit.Qubit):
+    """Apply the U3 gate on a single qubit.
+
+    Args:
+        theta (float): Rotation angle around the Y axis in radians.
+        phi (float): Rotation angle around the Z axis in radians.
+        lam (float): Rotation angle around the Z axis in radians.
+        qubit (qubit.Qubit): The qubit to apply the U3 gate to.
+    """
+    broadcast.u3(theta, phi, lam, ilist.IList([qubit]))
+
+
+@kernel
+def cz(control: qubit.Qubit, target: qubit.Qubit):
+    """Apply a controlled-Z gate on two qubits.
+
+    Args:
+        control (qubit.Qubit): The control qubit.
+        target (qubit.Qubit): The target qubit.
+    """
+    broadcast.cz(ilist.IList([control]), ilist.IList([target]))
+
+
+@kernel
+def cx(control: qubit.Qubit, target: qubit.Qubit):
+    """Apply a controlled-X gate on two qubits.
+
+    Args:
+        control (qubit.Qubit): The control qubit.
+        target (qubit.Qubit): The target qubit.
+    """
+    broadcast.cx(ilist.IList([control]), ilist.IList([target]))
+
+
+@kernel
+def cy(control: qubit.Qubit, targets: qubit.Qubit):
+    """Apply a controlled-Y gate on two qubits.
+
+    Args:
+        control (qubit.Qubit): The control qubit.
+        targets (qubit.Qubit): The target qubit.
+    """
+    broadcast.cy(ilist.IList([control]), ilist.IList([targets]))

bloqade/pyqrack/__init__.py

@@ -21,4 +21,5 @@ from .device import (
     StackMemorySimulator as StackMemorySimulator,
     DynamicMemorySimulator as DynamicMemorySimulator,
 )
+from .native import NativeMethods as NativeMethods
 from .target import PyQrack as PyQrack

bloqade/pyqrack/device.py

@@ -1,4 +1,4 @@
-from typing import Any, TypeVar, ParamSpec
+from typing import Any, TypeVar, ParamSpec, NamedTuple
 from dataclasses import field, dataclass
 
 import numpy as np
@@ -28,9 +28,99 @@ RetType = TypeVar("RetType")
 Params = ParamSpec("Params")
 
 
+class QuantumState(NamedTuple):
+    """
+    A representation of a quantum state as a density matrix, where the density matrix is
+    rho = sum_i eigenvalues[i] |eigenvectors[:,i]><eigenvectors[:,i]|.
+
+    This reprsentation is efficient for low-rank density matrices by only storing
+    the non-zero eigenvalues and corresponding eigenvectors of the density matrix.
+    For example, a pure state has only one non-zero eigenvalue equal to 1.0.
+
+    Endianness and qubit ordering of the state vector is consistent with Cirq, where
+    eigenvectors[0,0] corresponds to the amplitude of the |00..000> element of the zeroth eigenvector;
+    eigenvectors[1,0] corresponds to the amplitude of the |00..001> element of the zeroth eigenvector;
+    eigenvectors[3,0] corresponds to the amplitude of the |00..011> element of the zeroth eigenvector;
+    eigenvectors[-1,0] corresponds to the amplitude of the |11..111> element of the zeroth eigenvector.
+    A flip of the LAST bit |00..000><00..001| corresponds to applying a PauliX gate to the FIRST qubit.
+    A flip of the FIRST bit |00..000><10..000| corresponds to applying a PauliX gate to the LAST qubit.
+
+    Attributes:
+        eigenvalues (1d np.ndarray):
+            The non-zero eigenvalues of the density matrix.
+        eigenvectors (2d np.ndarray):
+            The corresponding eigenvectors of the density matrix,
+            where eigenvectors[:,i] is the i-th eigenvector.
+    Methods:
+        Not Implemented, pending https://github.com/QuEraComputing/bloqade-circuit/issues/447
+    """
+
+    eigenvalues: np.ndarray
+    eigenvectors: np.ndarray
+
+    def canonicalize(self, tol: float = 1e-12) -> "QuantumState":
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    def __add__(self, other: "QuantumState") -> "QuantumState":
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    def __mul__(self, scalar: float) -> "QuantumState":
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    @property
+    def dense(self) -> np.ndarray[tuple[int, int], np.complexfloating]:
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    def __matmul__(self, right: "cirq.Circuit") -> "QuantumState":  # noqa: F821
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    def expect(self, operator: Any) -> float:
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    def probability(self) -> np.ndarray[tuple[int], np.floating]:
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    def von_neumann_entropy(self) -> float:
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    @property
+    def qubit_basis(self) -> list[PyQrackQubit]:
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    def reduced_density_matrix(
+        self, qubits: list[PyQrackQubit], tol: float = 1e-12
+    ) -> "QuantumState":
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+    def overlap(self, other: "QuantumState") -> complex:
+        raise NotImplementedError(
+            "https://github.com/QuEraComputing/bloqade-circuit/issues/447"
+        )
+
+
 def _pyqrack_reduced_density_matrix(
     inds: tuple[int, ...], sim_reg: QrackSimulator, tol: float = 1e-12
-) -> "np.linalg._linalg.EighResult":
+) -> QuantumState:
     """
     Extract the reduced density matrix representing the state of a list
     of qubits from a PyQRack simulator register.
@@ -73,7 +163,7 @@ def _pyqrack_reduced_density_matrix(
     s = s[:, nonzero_inds]
     v = v[nonzero_inds] ** 2
     # Forge into the correct result type
-    result = np.linalg._linalg.EighResult(eigenvalues=v, eigenvectors=s)
+    result = QuantumState(eigenvalues=v, eigenvectors=s)
     return result
 
 
@@ -165,7 +255,7 @@ class PyQrackSimulatorBase(AbstractSimulatorDevice[PyQrackSimulatorTask]):
     @staticmethod
     def quantum_state(
         qubits: list[PyQrackQubit] | IList[PyQrackQubit, Any], tol: float = 1e-12
-    ) -> "np.linalg._linalg.EighResult":
+    ) -> "QuantumState":
         """
         Extract the reduced density matrix representing the state of a list
         of qubits from a PyQRack simulator register.
@@ -177,7 +267,7 @@ class PyQrackSimulatorBase(AbstractSimulatorDevice[PyQrackSimulatorTask]):
             An eigh result containing the eigenvalues and eigenvectors of the reduced density matrix.
         """
         if len(qubits) == 0:
-            return np.linalg._linalg.EighResult(
+            return QuantumState(
                 eigenvalues=np.array([]), eigenvectors=np.array([]).reshape(0, 0)
             )
         sim_reg = qubits[0].sim_reg

bloqade/pyqrack/native.py

@@ -0,0 +1,49 @@
+import math
+from typing import Any
+
+from kirin import interp
+from kirin.dialects import ilist
+
+from pyqrack import Pauli
+from bloqade.pyqrack import PyQrackQubit
+from bloqade.pyqrack.base import PyQrackInterpreter
+from bloqade.native.dialects import gates
+
+
+@gates.dialect.register(key="pyqrack")
+class NativeMethods(interp.MethodTable):
+
+    @interp.impl(gates.CZ)
+    def cz(self, _interp: PyQrackInterpreter, frame: interp.Frame, stmt: gates.CZ):
+        ctrls = frame.get_casted(stmt.ctrls, ilist.IList[PyQrackQubit, Any])
+        qargs = frame.get_casted(stmt.qargs, ilist.IList[PyQrackQubit, Any])
+
+        for ctrl, qarg in zip(ctrls, qargs):
+            if ctrl.is_active() and qarg.is_active():
+                ctrl.sim_reg.mcz([ctrl.addr], qarg.addr)
+
+        return ()
+
+    @interp.impl(gates.R)
+    def r(self, _interp: PyQrackInterpreter, frame: interp.Frame, stmt: gates.R):
+        inputs = frame.get_casted(stmt.inputs, ilist.IList[PyQrackQubit, Any])
+        rotation_angle = 2 * math.pi * frame.get_casted(stmt.rotation_angle, float)
+        axis_angle = 2 * math.pi * frame.get_casted(stmt.axis_angle, float)
+        for qubit in inputs:
+            if qubit.is_active():
+                qubit.sim_reg.r(Pauli.PauliZ, axis_angle, qubit.addr)
+                qubit.sim_reg.r(Pauli.PauliX, rotation_angle, qubit.addr)
+                qubit.sim_reg.r(Pauli.PauliZ, -axis_angle, qubit.addr)
+
+        return ()
+
+    @interp.impl(gates.Rz)
+    def rz(self, _interp: PyQrackInterpreter, frame: interp.Frame, stmt: gates.Rz):
+        inputs = frame.get_casted(stmt.inputs, ilist.IList[PyQrackQubit, Any])
+        rotation_angle = 2 * math.pi * frame.get_casted(stmt.rotation_angle, float)
+
+        for qubit in inputs:
+            if qubit.is_active():
+                qubit.sim_reg.r(Pauli.PauliZ, rotation_angle, qubit.addr)
+
+        return ()

bloqade/pyqrack/task.py

@@ -1,6 +1,10 @@
 from typing import TypeVar, ParamSpec, cast
+from collections import Counter
 from dataclasses import dataclass
 
+import numpy as np
+from kirin.dialects.ilist import IList
+
 from bloqade.task import AbstractSimulatorTask
 from bloqade.pyqrack.reg import QubitState, PyQrackQubit
 from bloqade.pyqrack.base import (
@@ -51,3 +55,93 @@ class PyQrackSimulatorTask(AbstractSimulatorTask[Param, RetType, MemoryType]):
         except AttributeError:
             Warning("Task has not been run, there are no qubits!")
             return []
+
+    def batch_run(self, shots: int = 1) -> dict[RetType, float]:
+        """
+        Repeatedly run the task to collect statistics on the shot outcomes.
+        The average is done over [shots] repetitions and thus is frequentist
+        and converges to exact only in the shots -> infinity limit.
+
+        Args:
+            shots (int):
+                the number of repetitions of the task
+        Returns:
+            dict[RetType, float]:
+                a dictionary mapping outcomes to their probabilities,
+                as estimated from counting the shot outcomes. RetType must be hashable.
+        """
+
+        results: list[RetType] = [self.run() for _ in range(shots)]
+
+        # Convert IList to tuple so that it is hashable by Counter
+        def convert(data):
+            if isinstance(data, (list, IList)):
+                return tuple(convert(item) for item in data)
+            return data
+
+        results = convert(results)
+
+        data = {
+            key: value / len(results) for key, value in Counter(results).items()
+        }  # Normalize to probabilities
+        return data
+
+    def batch_state(
+        self, shots: int = 1, qubit_map: None = None
+    ) -> "QuantumState":  # noqa: F821
+        """
+        Repeatedly run the task to extract the averaged quantum state.
+        The average is done over [shots] repetitions and thus is frequentist
+        and converges to exact only in the shots -> infinity limit.
+
+        Args:
+            shots (int):
+                the number of repetitions of the task
+            qubit_map (callable | None):
+                an optional callable that takes the output of self.run() and extract
+                the [returned] qubits to be used for the quantum state.
+                If None, all qubits in the simulator are used, in the order set by the simulator.
+                If callable, qubit_map must have the signature
+                > qubit_map(output:RetType) -> list[PyQrackQubit]
+                and the averaged state is
+                > quantum_state(qubit_map(self.run())).
+                If qubit_map is not None, self.run() must return qubit(s).
+                Two common patterns here are:
+                 > qubit_map = lambda qubits: qubits
+                for the case where self.run() returns a list of qubits, or
+                 > qubit_map = lambda qubit: [qubits]
+                for the case where self.run() returns a single qubit.
+        Returns:
+            QuantumState:
+                the averaged quantum state as a density matrix,
+                represented in its eigenbasis.
+        """
+        # Import here to avoid circular dependencies.
+        from bloqade.pyqrack.device import QuantumState, PyQrackSimulatorBase
+
+        states: list[QuantumState] = []
+        for _ in range(shots):
+            res = self.run()
+            if callable(qubit_map):
+                qbs = qubit_map(res)
+            else:
+                qbs = self.qubits()
+            states.append(PyQrackSimulatorBase.quantum_state(qbs))
+
+        state = QuantumState(
+            eigenvectors=np.concatenate(
+                [state.eigenvectors for state in states], axis=1
+            ),
+            eigenvalues=np.concatenate([state.eigenvalues for state in states], axis=0)
+            / len(states),
+        )
+
+        # Canonicalize the state by orthoganalizing the basis vectors.
+        tol = 1e-7
+        s, v, d = np.linalg.svd(
+            state.eigenvectors * np.sqrt(state.eigenvalues), full_matrices=False
+        )
+        mask = v > tol
+        v = v[mask] ** 2
+        s = s[:, mask]
+        return QuantumState(eigenvalues=v, eigenvectors=s)

bloqade/squin/cirq/emit/qubit.py

@@ -36,11 +36,9 @@ class EmitCirqQubitMethods(MethodTable):
         op = frame.get(stmt.operator)
         qbit_lists = [frame.get(qbit) for qbit in stmt.qubits]
 
-        cirq_ops = []
         for qbits in zip(*qbit_lists):
-            cirq_ops.extend(op.apply(qbits))
+            frame.circuit.append(op.apply(qbits))
 
-        frame.circuit.append(cirq.Moment(cirq_ops))
         return ()
 
     @impl(qubit.MeasureQubit)

{bloqade_circuit-0.7.4.dist-info → bloqade_circuit-0.7.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bloqade-circuit
-Version: 0.7.4
+Version: 0.7.6
 Summary: The software development toolkit for neutral atom arrays.
 Author-email: Roger-luo <rluo@quera.com>, kaihsin <khwu@quera.com>, weinbe58 <pweinberg@quera.com>, johnzl-777 <jlong@quera.com>
 License-File: LICENSE

{bloqade_circuit-0.7.4.dist-info → bloqade_circuit-0.7.6.dist-info}/RECORD

@@ -21,12 +21,23 @@ bloqade/cirq_utils/noise/_two_zone_utils.py,sha256=iq4nwdJQITFlGB61wfrV7vyPA2194
 bloqade/cirq_utils/noise/conflict_graph.py,sha256=ZUwPWTknrb6SgtZUVPeICn3YA-nUeWQJDuKKX5jL9tE,7179
 bloqade/cirq_utils/noise/model.py,sha256=8qovvB50oHzDszXkuMAs2I9BQV3eS1IU2D7Wux_dsGE,18459
 bloqade/cirq_utils/noise/transform.py,sha256=pauFnOKbk2QjxeyXEV_x2zyRGypr5wiQ6ySirU7C2zg,2278
-bloqade/pyqrack/__init__.py,sha256=lonTS-luJkTVujCCtgdZRC12V7FQdoFcozAI-byXwN0,810
+bloqade/native/__init__.py,sha256=MtQZmq7KqNGkjxsO_22DpY-13X23Ovw8JKs7lfuv55w,443
+bloqade/native/_prelude.py,sha256=znle4mSTLZtRMCmmiZ-bCPfHqNdgVxsKuY64Fgj2kko,1258
+bloqade/native/dialects/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bloqade/native/dialects/gates/__init__.py,sha256=YX5h-5Agl_x5K0KD-QnuloNV8oeIFXnsGT4ouMpabCU,138
+bloqade/native/dialects/gates/_dialect.py,sha256=D5vU_JAE2ROT5jMMth3f9SPjL9vfvCVjzEfLJzvAImI,61
+bloqade/native/dialects/gates/_interface.py,sha256=2Tie1C-k9GrrSJEsrQ3IIO_3yYJXsyoD6pBZmm0cyp4,537
+bloqade/native/dialects/gates/stmts.py,sha256=qBziqe1b39ltwcFUEpBZ2RdzL4DTiesRs__HIjbA5N4,1016
+bloqade/native/stdlib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+bloqade/native/stdlib/broadcast.py,sha256=0DEK6pRcz0IG5BZHt0sV1wZpDvlJgvw4-Qwd1xYLSYM,6141
+bloqade/native/stdlib/simple.py,sha256=mK_p9x_TUr4g-q-Q4LfbVBtXxBS21XPh1meyCC4to_k,5416
+bloqade/pyqrack/__init__.py,sha256=OV8-2fw44gP_JgY8mAUiwISO_qYxS-t0fKsbuUB-r9Y,861
 bloqade/pyqrack/base.py,sha256=g0GRlEgyJ_P8z-lR8RK2CAuRTj6KPfglKX0iwrgg4DM,4408
-bloqade/pyqrack/device.py,sha256=-zlr1lSzOvcj5l28nnevy1oMYct79DTOdLYyGaT2Yco,11567
+bloqade/pyqrack/device.py,sha256=ruseNWt3scvjw0KyHZjwnM6Z6lFjROgmW3Zdi1lucgQ,15195
+bloqade/pyqrack/native.py,sha256=ErbVQCatn_JT3Ej-iQzMMfb_q50JF_K1Iv1vRYvu5VA,1857
 bloqade/pyqrack/reg.py,sha256=uTL07CT1R0xUsInLmwU9YuuNdV6lV0lCs1zhdUz1qIs,1660
 bloqade/pyqrack/target.py,sha256=c78VtLWAiDNp_0sXwvVzhaEoeFsr1fUVsupxWuo6p3s,3661
-bloqade/pyqrack/task.py,sha256=hzMueE03C-MR_-ti_5Z0lqFZcHDVCxmQHFaX9tPNm2o,1522
+bloqade/pyqrack/task.py,sha256=nRmI3tM_y4134_Uld-D9xwNXG9ie-OHNAqwRT5Ss-Y0,5266
 bloqade/pyqrack/noise/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 bloqade/pyqrack/noise/native.py,sha256=BmRlRzqCzvSgfSsDlIjuB8L0gx4uaPqmPyfJ5-Ole6M,2751
 bloqade/pyqrack/qasm2/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -145,7 +156,7 @@ bloqade/squin/cirq/lowering.py,sha256=trrcBDQQzy2r3JVKb2x6i2BdzBMYkjQNBzH_Va6JOu
 bloqade/squin/cirq/emit/emit_circuit.py,sha256=JVFXiaSB7A9kamRwCjLqs03Sel4PVCBT-6kRNRWX-jo,4393
 bloqade/squin/cirq/emit/noise.py,sha256=rESjGC_66s2Y4FwwYda4rY3mYHYjbqLlKE_vnqpZDYI,1534
 bloqade/squin/cirq/emit/op.py,sha256=u1FoOlbz7NNXPlU-hMbnX_yRB5suKbsEm2jF3Ts4f68,5764
-bloqade/squin/cirq/emit/qubit.py,sha256=FManFeDyi1l4Sd07IeD7Lb4xR-7jy5Nwe-quQHI1yPw,1984
+bloqade/squin/cirq/emit/qubit.py,sha256=9aDM0USxcekzXx9--WJB5r96JQ1iEmeNJ9h8rdZ8QhM,1915
 bloqade/squin/cirq/emit/runtime.py,sha256=dH7JSMt2mALPhVFjmZETQzvnTUQ3BFY5poe0YZpM5vQ,6819
 bloqade/squin/noise/__init__.py,sha256=JFJ4kmEeWt6bJ2xx3yA5ek-NEbj8ilYWf6mepgbJLOo,338
 bloqade/squin/noise/_dialect.py,sha256=2IR98J-lXm5Y3srP9g-FD4JC-qTq2seureM6mKKq1xg,63
@@ -232,7 +243,7 @@ bloqade/visual/animation/runtime/atoms.py,sha256=EmjxhujLiHHPS_HtH_B-7TiqeHgvW5u
 bloqade/visual/animation/runtime/ppoly.py,sha256=JB9IP53N1w6adBJEue6J5Nmj818Id9JvrlgrmiQTU1I,1385
 bloqade/visual/animation/runtime/qpustate.py,sha256=rlmxQeJSvaohXrTpXQL5y-NJcpvfW33xPaYM1slv7cc,4270
 bloqade/visual/animation/runtime/utils.py,sha256=ju9IzOWX-vKwfpqUjlUKu3Ssr_UFPFFq-tzH_Nqyo_c,1212
-bloqade_circuit-0.7.4.dist-info/METADATA,sha256=pvbmLzuKl7jWDD9tjP6Z5VckwUv4sqUb00ArREsgXwo,3850
-bloqade_circuit-0.7.4.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-bloqade_circuit-0.7.4.dist-info/licenses/LICENSE,sha256=S5GIJwR6QCixPA9wryYb44ZEek0Nz4rt_zLUqP05UbU,13160
-bloqade_circuit-0.7.4.dist-info/RECORD,,
+bloqade_circuit-0.7.6.dist-info/METADATA,sha256=KBi8hGUeoSp9gEeZUOTjO7QQxUaZHk5pYH7gWCUqB3E,3850
+bloqade_circuit-0.7.6.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+bloqade_circuit-0.7.6.dist-info/licenses/LICENSE,sha256=S5GIJwR6QCixPA9wryYb44ZEek0Nz4rt_zLUqP05UbU,13160
+bloqade_circuit-0.7.6.dist-info/RECORD,,