Add support for StatePrep and BasisState in QJIT with PLxPR program capture (#1631)

**Context:** Currently, when you try to execute a circuit containing
`qml.StatePrep()` through the QJIT pipeline with PLxPR program capture
enabled it results in an error, since the operation is bound to the
`qinst_p` primitive rather than the `set_state_p` primitive. `qinst_p`
primitives expect input parameters to be of type `float64`, but
`StatePrep` can have complex input values, hence the error (see issue
#1630 for details).

**Description of the Change:** Binds `qml.StatePrep` operations to the
`set_state_p` primitive to correctly support StatePrep in QJIT with
PLxPR program capture. This PR also adds support for `qml.BasisState()`
in QJIT with PLxPR capture since it is closely related to StatePrep.

**Related GitHub Issues:** Fixes #1630


[[sc-88638](https://app.shortcut.com/xanaduai/story/88638/make-qml-stateprep-qjit-compatible-with-plxpr-program-capture),
[sc-88653](https://app.shortcut.com/xanaduai/story/88653/stateprep-not-supported-in-qjit-with-plxpr-program-capture-enabled)]

Author: joeycarter

doc/releases/changelog-dev.md

@@ -19,6 +19,10 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Catalyst now correctly supports `qml.StatePrep()` and `qml.BasisState()` operations in the
+  experimental PennyLane program-capture pipeline.
+  [(#1631)](https://github.com/PennyLaneAI/catalyst/pull/1631)
+
 <h3>Internal changes ⚙️</h3>
 
 <h3>Documentation 📝</h3>

frontend/catalyst/from_plxpr.py

@@ -20,6 +20,7 @@
 
 import jax
 import jax.core
+import jax.numpy as jnp
 import pennylane as qml
 from jax.extend.linear_util import wrap_init
 from jax.interpreters.partial_eval import convert_constvars_jaxpr
@@ -65,6 +66,8 @@
     quantum_kernel_p,
     qunitary_p,
     sample_p,
+    set_basis_state_p,
+    set_state_p,
     state_p,
     var_p,
     while_p,
@@ -423,6 +426,37 @@ def handle_global_phase(self, phase, *wires, n_wires):
     gphase_p.bind(phase, ctrl_len=0, adjoint=False)
 
 
+@QFuncPlxprInterpreter.register_primitive(qml.BasisState._primitive)
+def handle_basis_state(self, *invals, n_wires):
+    """Handle the conversion from plxpr to Catalyst jaxpr for the BasisState primitive"""
+    state_inval = invals[0]
+    wires_inval = invals[1:]
+
+    state = jax.lax.convert_element_type(state_inval, jnp.dtype(jnp.bool))
+    wires = [self.get_wire(w) for w in wires_inval]
+    out_wires = set_basis_state_p.bind(*wires, state)
+
+    for wire_values, new_wire in zip(wires_inval, out_wires):
+        self.wire_map[wire_values] = new_wire
+
+
+# pylint: disable=unused-argument
+@QFuncPlxprInterpreter.register_primitive(qml.StatePrep._primitive)
+def handle_state_prep(self, *invals, n_wires, **kwargs):
+    """Handle the conversion from plxpr to Catalyst jaxpr for the StatePrep primitive"""
+    state_inval = invals[0]
+    wires_inval = invals[1:]
+
+    # jnp.complex128 is the top element in the type promotion lattice so it is ok to do this:
+    # https://jax.readthedocs.io/en/latest/type_promotion.html
+    state = jax.lax.convert_element_type(state_inval, jnp.dtype(jnp.complex128))
+    wires = [self.get_wire(w) for w in wires_inval]
+    out_wires = set_state_p.bind(*wires, state)
+
+    for wire_values, new_wire in zip(wires_inval, out_wires):
+        self.wire_map[wire_values] = new_wire
+
+
 # pylint: disable=unused-argument, too-many-arguments
 @QFuncPlxprInterpreter.register_primitive(plxpr_cond_prim)
 def handle_cond(self, *plxpr_invals, jaxpr_branches, consts_slices, args_slice):

frontend/test/pytest/test_capture_integration.py

@@ -177,6 +177,57 @@ def g(x):
 
         assert jnp.allclose(actual, desired)
 
+    @pytest.mark.parametrize(
+        "n_wires, basis_state",
+        [
+            (1, jnp.array([0])),
+            (1, jnp.array([1])),
+            (2, jnp.array([0, 0])),
+            (2, jnp.array([0, 1])),
+            (2, jnp.array([1, 0])),
+            (2, jnp.array([1, 1])),
+        ],
+    )
+    def test_basis_state(self, backend, n_wires, basis_state):
+        """Test the integration for a circuit with BasisState."""
+        dev = qml.device(backend, wires=n_wires)
+
+        @qml.qnode(dev)
+        def circuit(_basis_state):
+            qml.BasisState(_basis_state, wires=list(range(n_wires)))
+            return qml.state()
+
+        desired = circuit(basis_state)
+        actual = qjit(circuit, experimental_capture=True)(basis_state)
+
+        assert jnp.allclose(actual, desired)
+
+    @pytest.mark.parametrize(
+        "n_wires, init_state",
+        [
+            (1, jnp.array([1.0, 0.0], dtype=jnp.complex128)),
+            (1, jnp.array([0.0, 1.0], dtype=jnp.complex128)),
+            (1, jnp.array([1.0, 1.0], dtype=jnp.complex128) / jnp.sqrt(2.0)),
+            (1, jnp.array([0.0, 1.0], dtype=jnp.float64)),
+            (1, jnp.array([0, 1], dtype=jnp.int64)),
+            (2, jnp.array([1.0, 0.0, 0.0, 0.0], dtype=jnp.complex128)),
+            (2, jnp.array([0.0, 1.0, 0.0, 0.0], dtype=jnp.complex128)),
+        ],
+    )
+    def test_state_prep(self, backend, n_wires, init_state):
+        """Test the integration for a circuit with StatePrep."""
+        dev = qml.device(backend, wires=n_wires)
+
+        @qml.qnode(dev)
+        def circuit(init_state):
+            qml.StatePrep(init_state, wires=list(range(n_wires)))
+            return qml.state()
+
+        desired = circuit(init_state)
+        actual = qjit(circuit, experimental_capture=True)(init_state)
+
+        assert jnp.allclose(actual, desired)
+
     @pytest.mark.xfail(reason="Adjoint not supported.")
     @pytest.mark.parametrize("theta, val", [(jnp.pi, 0), (-100.0, 1)])
     def test_adjoint(self, backend, theta, val):

frontend/test/pytest/test_from_plxpr.py

@@ -48,17 +48,40 @@ def capture(self, args):
     return JAXPRRunner(fn=lambda: None, compile_options=catalyst.CompileOptions())
 
 
-def compare_call_jaxprs(jaxpr1, jaxpr2, skip_eqns=()):
+def compare_call_jaxprs(jaxpr1, jaxpr2, skip_eqns=(), ignore_order=False):
     """Compares two call jaxprs and validates that they are essentially equal."""
     for inv1, inv2 in zip(jaxpr1.invars, jaxpr2.invars):
         assert inv1.aval == inv2.aval, f"{inv1.aval}, {inv2.aval}"
     for ov1, ov2 in zip(jaxpr1.outvars, jaxpr2.outvars):
         assert ov1.aval == ov2.aval
-    assert len(jaxpr1.eqns) == len(jaxpr2.eqns)
-
-    for i, (eqn1, eqn2) in enumerate(zip(jaxpr1.eqns, jaxpr2.eqns)):
-        if i not in skip_eqns:
-            compare_eqns(eqn1, eqn2)
+    assert len(jaxpr1.eqns) == len(
+        jaxpr2.eqns
+    ), f"Number of equations differ: {len(jaxpr1.eqns)} vs {len(jaxpr2.eqns)}"
+
+    if not ignore_order:
+        # Assert that equations in both jaxprs are equivalent and in same order
+        for i, (eqn1, eqn2) in enumerate(zip(jaxpr1.eqns, jaxpr2.eqns)):
+            if i not in skip_eqns:
+                compare_eqns(eqn1, eqn2)
+
+    else:
+        # Assert that equations in both jaxprs are equivalent but in any order
+        eqns1 = [eqn for i, eqn in enumerate(jaxpr1.eqns) if i not in skip_eqns]
+        eqns2 = [eqn for i, eqn in enumerate(jaxpr2.eqns) if i not in skip_eqns]
+
+        for eqn1 in eqns1:
+            found_match = False
+            for i, eqn2 in enumerate(eqns2):
+                try:
+                    compare_eqns(eqn1, eqn2)
+                    # Remove the matched equation to prevent double-matching
+                    eqns2.pop(i)
+                    found_match = True
+                    break  # Exit inner loop after finding a match
+                except AssertionError:
+                    pass  # Continue to the next equation in eqns2
+            if not found_match:
+                raise AssertionError(f"No matching equation found for: {eqn1}")
 
 
 def compare_eqns(eqn1, eqn2):
@@ -428,6 +451,73 @@ def circuit():
 
         compare_call_jaxprs(call_jaxpr_pl, call_jaxpr_c)
 
+    def test_basis_state(self, disable_capture):
+        """Test comparison and execution of a jaxpr containing BasisState."""
+        dev = qml.device("lightning.qubit", wires=2)
+
+        @qml.qnode(dev)
+        def circuit(_basis_state):
+            qml.BasisState(_basis_state, wires=[0, 1])
+            return qml.state()
+
+        basis_state = np.array([1, 1])
+        expected_state_vector = np.array([0, 0, 0, 1], dtype=np.complex128)
+
+        qml.capture.enable()
+        plxpr = jax.make_jaxpr(circuit)(basis_state)
+        converted = from_plxpr(plxpr)(basis_state)
+        qml.capture.disable()
+
+        assert converted.eqns[0].primitive == catalyst.jax_primitives.quantum_kernel_p
+        assert converted.eqns[0].params["qnode"] is circuit
+
+        catalyst_res = catalyst_execute_jaxpr(converted)(basis_state)
+        assert len(catalyst_res) == 1
+        assert qml.math.allclose(catalyst_res[0], expected_state_vector)
+
+        qjit_obj = qjit(circuit)
+        qjit_obj(basis_state)
+        catalxpr = qjit_obj.jaxpr
+        call_jaxpr_pl = get_call_jaxpr(converted)
+        call_jaxpr_c = get_call_jaxpr(catalxpr)
+
+        # Ignore ordering of eqns when comparing jaxpr since Catalyst performs sorting
+        compare_call_jaxprs(call_jaxpr_pl, call_jaxpr_c, ignore_order=True)
+
+    def test_state_prep(self, disable_capture):
+        """Test comparison and execution of a jaxpr containing StatePrep."""
+        dev = qml.device("lightning.qubit", wires=1)
+
+        @qml.qnode(dev)
+        def circuit(_init_state):
+            # NOTE: Require validate_norm=False here otherwise Catalyst jaxpr contains
+            # unused function that computes norm
+            qml.StatePrep(_init_state, wires=0, validate_norm=False)
+            return qml.state()
+
+        init_state = np.array([1, 1], dtype=np.complex128) / np.sqrt(2)
+
+        qml.capture.enable()
+        plxpr = jax.make_jaxpr(circuit)(init_state)
+        converted = from_plxpr(plxpr)(init_state)
+        qml.capture.disable()
+
+        assert converted.eqns[0].primitive == catalyst.jax_primitives.quantum_kernel_p
+        assert converted.eqns[0].params["qnode"] is circuit
+
+        catalyst_res = catalyst_execute_jaxpr(converted)(init_state)
+        assert len(catalyst_res) == 1
+        assert qml.math.allclose(catalyst_res[0], init_state)
+
+        qjit_obj = qjit(circuit)
+        qjit_obj(init_state)
+        catalxpr = qjit_obj.jaxpr
+        call_jaxpr_pl = get_call_jaxpr(converted)
+        call_jaxpr_c = get_call_jaxpr(catalxpr)
+
+        # Ignore ordering of eqns when comparing jaxpr since Catalyst performs sorting
+        compare_call_jaxprs(call_jaxpr_pl, call_jaxpr_c, ignore_order=True)
+
     def test_multiple_measurements(self, disable_capture):
         """Test that we can convert a circuit with multiple measurement returns."""