[BUG] `decompose_mcx` is not JIT-compatible

[dwierichs]

Expected behavior

MultiControlledX can be decomposed with any number of work wires and QJIT.

Actual behavior

• When there is a single work wire and QJIT is used, decompose_mcx uses adjoint with lazy=False on a callable, which is not supported by Catalyst.
• When there are two work wires and QJIT is used, decompose_mcx uses a qml.for_loop but casts input wires to a list before using them, raising an error.
• When there are more than two work wires, the same problem as for two occurs, caused by a different for-loop though.

Additional information

No response

Source code

import pennylane as qp
from pennylane.ops.op_math.controlled_decompositions import decompose_mcx

qp.decomposition.disable_graph()

qp.qjit(decompose_mcx)([0, 1, 2, 3, 4], 5, [6])

Tracebacks

...
File ~/venvs/dev/lib/python3.13/site-packages/catalyst/api_extensions/quantum_operators.py:254, in adjoint(f, lazy)
    188 def adjoint(f: Union[Callable, Operator], lazy=True) -> Union[Callable, Operator]:
    189     """A :func:`~.qjit` compatible adjoint transformer for PennyLane/Catalyst.
    190 
    191     Returns a quantum function or operator that applies the adjoint of the
   (...)    251     [1.00000000e+00 7.39557099e-32]
    252     """
--> 254     adj = AdjointCallable(f, lazy=lazy)
    256     # Return an instantiated version of the Adjoint class if we receive an operator instance.
    257     return adj() if isinstance(f, Operator) else adj

File ~/venvs/dev/lib/python3.13/site-packages/catalyst/api_extensions/quantum_operators.py:399, in AdjointCallable.__init__(self, target, lazy)
    392     raise ValueError(f"Expected a callable or a qml.Operator, not {target}")
    394 if not self.lazy and not self.single_op:
    395     # Supporting this for a qfunc is technically possible by invoking the decomposition on
    396     # HybridAdjoint with laza=False, however one would need to outline the classical jaxpr
    397     # into the outer scope, and possibly re-mapping tracers in the quantum operators,
    398     # in order to avoid escaped tracers which indicate an invalid program.
--> 399     raise ValueError(
    400         "Eagerly computing the adjoint (lazy=False) is only supported on single operators."
    401     )

ValueError: Eagerly computing the adjoint (lazy=False) is only supported on single operators.

System information

pl dev

Existing GitHub issues

• I have searched existing GitHub issues to make sure the issue does not already exist.