Making arithmetic templates QJIT compatible

doc/releases/changelog-dev.md

@@ -39,6 +39,9 @@
 * The `Hermitian` operator now has a `compute_sparse_matrix` implementation.
   [(#6225)](https://github.com/PennyLaneAI/pennylane/pull/6225)
 
+* The quantum arithmetic templates are now QJIT compatible.
+  [(#6307)](https://github.com/PennyLaneAI/pennylane/pull/6307)
+
 <h4>Capturing and representing hybrid programs</h4>
 
 * Differentiation of hybrid programs via `qml.grad` and `qml.jacobian` can now be captured

pennylane/templates/subroutines/multiplier.py

@@ -228,5 +228,9 @@ def compute_decomposition(k, x_wires, mod, work_wires):  # pylint: disable=argum
         for x_wire, aux_wire in zip(x_wires, wires_aux_swap):
             op_list.append(qml.SWAP(wires=[x_wire, aux_wire]))
         inv_k = pow(k, -1, mod)
-        op_list.extend(qml.adjoint(_mul_out_k_mod)(inv_k, x_wires, mod, work_wire_aux, wires_aux))
+
+        # Adjoint is not iterable in QJIT
+        for op in reversed(_mul_out_k_mod(inv_k, x_wires, mod, work_wire_aux, wires_aux)):
+            op_list.append(qml.adjoint(op))
+
         return op_list

pennylane/templates/subroutines/phase_adder.py

@@ -131,20 +131,25 @@ def __init__(
 
         num_work_wires = 0 if work_wire is None else len(work_wire)
 
-        if mod is None:
-            mod = 2 ** len(x_wires)
-        elif mod != 2 ** len(x_wires) and num_work_wires != 1:
-            raise ValueError(f"If mod is not 2^{len(x_wires)}, one work wire should be provided.")
-        if not isinstance(k, int) or not isinstance(mod, int):
-            raise ValueError("Both k and mod must be integers")
-        if mod > 2 ** len(x_wires):
-            raise ValueError(
-                "PhaseAdder must have enough x_wires to represent mod. The maximum mod "
-                f"with len(x_wires)={len(x_wires)} is {2 ** len(x_wires)}, but received {mod}."
-            )
-        if work_wire is not None:
-            if any(wire in work_wire for wire in x_wires):
-                raise ValueError("None of the wires in work_wire should be included in x_wires.")
+        if not qml.math.is_abstract(mod):
+            if mod is None:
+                mod = 2 ** len(x_wires)
+            elif mod != 2 ** len(x_wires) and num_work_wires != 1:
+                raise ValueError(
+                    f"If mod is not 2^{len(x_wires)}, one work wire should be provided."
+                )
+            if not isinstance(k, int) or not isinstance(mod, int):
+                raise ValueError("Both k and mod must be integers")
+            if mod > 2 ** len(x_wires):
+                raise ValueError(
+                    "PhaseAdder must have enough x_wires to represent mod. The maximum mod "
+                    f"with len(x_wires)={len(x_wires)} is {2 ** len(x_wires)}, but received {mod}."
+                )
+            if work_wire is not None:
+                if any(wire in work_wire for wire in x_wires):
+                    raise ValueError(
+                        "None of the wires in work_wire should be included in x_wires."
+                    )
 
         self.hyperparameters["k"] = k % mod
         self.hyperparameters["mod"] = mod
@@ -216,12 +221,20 @@ def compute_decomposition(k, x_wires, mod, work_wire):  # pylint: disable=argume
         else:
             aux_k = x_wires[0]
             op_list.extend(_add_k_fourier(k, x_wires))
-            op_list.extend(qml.adjoint(_add_k_fourier)(mod, x_wires))
+
+            # Adjoint is not iterable in QJIT
+            for op in reversed(_add_k_fourier(mod, x_wires)):
+                op_list.append(qml.adjoint(op))
+
             op_list.append(qml.adjoint(qml.QFT)(wires=x_wires))
             op_list.append(qml.ctrl(qml.PauliX(work_wire), control=aux_k, control_values=1))
             op_list.append(qml.QFT(wires=x_wires))
             op_list.extend(qml.ctrl(op, control=work_wire) for op in _add_k_fourier(mod, x_wires))
-            op_list.extend(qml.adjoint(_add_k_fourier)(k, x_wires))
+
+            # Adjoint is not iterable in QJIT
+            for op in reversed(_add_k_fourier(k, x_wires)):
+                op_list.append(qml.adjoint(op))
+
             op_list.append(qml.adjoint(qml.QFT)(wires=x_wires))
             op_list.append(qml.ctrl(qml.PauliX(work_wire), control=aux_k, control_values=0))
             op_list.append(qml.QFT(wires=x_wires))