[BUG] ControlledQubitUnitary as incorrect wires when using with qml.capture.enable()

[mudit2812]

Expected behavior

I expect the same class signature usage to work for ControlledQubitUnitary with program capture enabled and disabled.

Actual behavior

Using the calling signature ControlledQubitUnitary(U_mat, control_wires, target_wires) works only with program capture disabled (this is the documented signature). For ControlledQubitUnitary to be usable with program capture, the signature that needs to be used is ControlledQubitUnitary(U_mat, all_wires).

Additional information

No response

Source code

>>> import pennylane as qml
>>> qml.capture.disable()
>>> print(qml.ControlledQubitUnitary([[0, 1], [1, 0]], 0, 1))
Controlled(QubitUnitary(array([[0, 1],
       [1, 0]]), wires=[1]), control_wires=[0])
>>> qml.capture.enable()
>>> print(qml.ControlledQubitUnitary([[0, 1], [1, 0]], [0, 1]))
Controlled(QubitUnitary(array([[0, 1],
       [1, 0]]), wires=[1]), control_wires=[0])
>>> print(qml.ControlledQubitUnitary([[0, 1], [1, 0]], 0, 1))
Error (traceback below)

Tracebacks

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/mudit.pandey/repos/pennylane/pennylane/capture/capture_meta.py", line 88, in __call__
    return cls._primitive_bind_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mudit.pandey/repos/pennylane/pennylane/ops/op_math/controlled.py", line 487, in _primitive_bind_call
    return cls._primitive.bind(
           ^^^^^^^^^^^^^^^^^^^^
  File "/Users/mudit.pandey/.pyenv/versions/pennylane/lib/python3.11/site-packages/jax/_src/core.py", line 387, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mudit.pandey/repos/pennylane/pennylane/capture/capture_diff.py", line 50, in bind_with_trace
    return super().bind_with_trace(trace, args, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mudit.pandey/.pyenv/versions/pennylane/lib/python3.11/site-packages/jax/_src/core.py", line 391, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mudit.pandey/.pyenv/versions/pennylane/lib/python3.11/site-packages/jax/_src/core.py", line 879, in process_primitive
    return primitive.impl(*tracers, **params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mudit.pandey/repos/pennylane/pennylane/capture/capture_operators.py", line 111, in _
    return type.__call__(operator_type, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mudit.pandey/repos/pennylane/pennylane/ops/op_math/controlled_ops.py", line 139, in __init__
    base = QubitUnitary(base, wires=wires, unitary_check=unitary_check)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mudit.pandey/repos/pennylane/pennylane/capture/capture_meta.py", line 88, in __call__
    return cls._primitive_bind_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mudit.pandey/repos/pennylane/pennylane/operation.py", line 712, in _primitive_bind_call
    return cls._primitive.bind(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mudit.pandey/.pyenv/versions/pennylane/lib/python3.11/site-packages/jax/_src/core.py", line 387, in bind
    return self.bind_with_trace(find_top_trace(args), args, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mudit.pandey/repos/pennylane/pennylane/capture/capture_diff.py", line 50, in bind_with_trace
    return super().bind_with_trace(trace, args, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mudit.pandey/.pyenv/versions/pennylane/lib/python3.11/site-packages/jax/_src/core.py", line 391, in bind_with_trace
    out = trace.process_primitive(self, map(trace.full_raise, args), params)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mudit.pandey/.pyenv/versions/pennylane/lib/python3.11/site-packages/jax/_src/core.py", line 879, in process_primitive
    return primitive.impl(*tracers, **params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mudit.pandey/repos/pennylane/pennylane/capture/capture_operators.py", line 118, in _
    wires = tuple(w if qml.math.is_abstract(w) else int(w) for w in args[split:])
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/mudit.pandey/repos/pennylane/pennylane/capture/capture_operators.py", line 118, in <genexpr>
    wires = tuple(w if qml.math.is_abstract(w) else int(w) for w in args[split:])
                                                    ^^^^^^
TypeError: int() argument must be a string, a bytes-like object or a real number, not 'NoneType'

System information

Name: PennyLane
Version: 0.40.0.dev26
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /Users/mudit.pandey/repos/pennylane
Requires: appdirs, autograd, autoray, cachetools, diastatic-malt, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, toml, typing_extensions
Required-by: PennyLane_Lightning

Platform info:           macOS-14.7.1-arm64-arm-64bit
Python version:          3.11.10
Numpy version:           2.0.0
Scipy version:           1.14.1
Installed devices:
- default.clifford (PennyLane-0.40.0.dev26)
- default.gaussian (PennyLane-0.40.0.dev26)
- default.mixed (PennyLane-0.40.0.dev26)
- default.qubit (PennyLane-0.40.0.dev26)
- default.qutrit (PennyLane-0.40.0.dev26)
- default.qutrit.mixed (PennyLane-0.40.0.dev26)
- default.tensor (PennyLane-0.40.0.dev26)
- null.qubit (PennyLane-0.40.0.dev26)
- reference.qubit (PennyLane-0.40.0.dev26)
- lightning.qubit (PennyLane_Lightning-0.39.0)

Existing GitHub issues

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