Repository: PennyLaneAI/pennylane

Issue:
Broadcast Expand Prevents Torchlayer from detecting Inputs Argument
### Expected behavior

The torchlayer should be created successfully without any errors. This works fine if the broadcast expand function decorator "@qml.transforms.broadcast_expand" is commented out from the minimal code provided.



### Actual behavior

Torchlayer fails to detect the inputs argument to the circuit function under the qnode and broadcast expand decorators.

### Additional information

The versions of the relevant packages are as below:
# Name                    Version                   Build  Channel
pennylane                 0.36.0                   pypi_0    pypi
python                    3.11.9               h955ad1f_0
torch                     2.3.1+cu121              pypi_0    pypi

### Source code

```shell
import pennylane as qml

dev = qml.device("default.qubit", wires = 1)

@qml.qnode(dev)
@qml.transforms.broadcast_expand
def circuit(inputs):
    return qml.probs(wires = 1)

qlayer = qml.qnn.TorchLayer(circuit, {})
```


### Tracebacks

```shell
line 10, in <module>
    qlayer = qml.qnn.TorchLayer(circuit, {})
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
pennylane/qnn/torch.py", line 351, in __init__
    self._signature_validation(qnode, weight_shapes)
pennylane/qnn/torch.py", line 364, in _signature_validation
    raise TypeError(
TypeError: QNode must include an argument with name inputs for inputting data
```


### System information

```shell
Name: PennyLane
Version: 0.36.0
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
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane-Catalyst, pennylane-qulacs, PennyLane_Lightning, PennyLane_Lightning_GPU

Platform info:           Linux-5.15.146.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
Python version:          3.11.9
Numpy version:           1.26.4
Scipy version:           1.13.1
Installed devices:
- lightning.qubit (PennyLane_Lightning-0.36.0)
- qulacs.simulator (pennylane-qulacs-0.36.0)
- nvidia.custatevec (PennyLane-Catalyst-0.6.0)
- nvidia.cutensornet (PennyLane-Catalyst-0.6.0)
- oqc.cloud (PennyLane-Catalyst-0.6.0)
- softwareq.qpp (PennyLane-Catalyst-0.6.0)
- default.clifford (PennyLane-0.36.0)
- default.gaussian (PennyLane-0.36.0)
- default.mixed (PennyLane-0.36.0)
- default.qubit (PennyLane-0.36.0)
- default.qubit.autograd (PennyLane-0.36.0)
- default.qubit.jax (PennyLane-0.36.0)
- default.qubit.legacy (PennyLane-0.36.0)
- default.qubit.tf (PennyLane-0.36.0)
- default.qubit.torch (PennyLane-0.36.0)
- default.qutrit (PennyLane-0.36.0)
- default.qutrit.mixed (PennyLane-0.36.0)
- null.qubit (PennyLane-0.36.0)
- lightning.gpu (PennyLane_Lightning_GPU-0.36.0)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.

Discussion:
albi3ro: > Thanks! I cannot find `inspect` among the required packages, but looks like is already used everywhere in the source code

`inspect` is a built in python package.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 11e431fa3..19dcc314a 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -264,6 +264,9 @@
 * `qml.qchem.molecular_dipole` function is added for calculating the dipole operator using "dhf" and "openfermion" backends.
   [(#5764)](https://github.com/PennyLaneAI/pennylane/pull/5764)
 
+* Transforms applied to callables now use `functools.wraps` to preserve the docstring and call signature of the original function.
+  [(#5857)](https://github.com/PennyLaneAI/pennylane/pull/5857)
+
 <h4>Community contributions 🥳</h4>
 
 * Implemented kwargs (`check_interface`, `check_trainability`, `rtol` and `atol`) support in `qml.equal` for the operators `Pow`, `Adjoint`, `Exp`, and `SProd`.
diff --git a/pennylane/transforms/core/transform_dispatcher.py b/pennylane/transforms/core/transform_dispatcher.py
index 803136924..7591c9d22 100644
--- a/pennylane/transforms/core/transform_dispatcher.py
+++ b/pennylane/transforms/core/transform_dispatcher.py
@@ -235,6 +235,7 @@ class TransformDispatcher:
     def _qfunc_transform(self, qfunc, targs, tkwargs):
         """Apply the transform on a quantum function."""
 
+        @functools.wraps(qfunc)
         def qfunc_transformed(*args, **kwargs):
             with qml.queuing.AnnotatedQueue() as q:
                 qfunc_output = qfunc(*args, **kwargs)


Test patch:
diff --git a/tests/transforms/core/test_transform_dispatcher.py b/tests/transforms/core/test_transform_dispatcher.py
index 21acd922b..64a008dec 100644
--- a/tests/transforms/core/test_transform_dispatcher.py
+++ b/tests/transforms/core/test_transform_dispatcher.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """Unit and integration tests for the transform dispatcher."""
+import inspect
 from functools import partial
 from typing import Callable, Sequence
 
@@ -31,7 +32,7 @@ with qml.tape.QuantumTape() as tape_circuit:
     qml.expval(qml.PauliZ(wires=0))
 
 
-def qfunc_circuit(a):
+def qfunc_circuit(a: qml.typing.TensorLike):
     """Qfunc circuit/"""
     qml.Hadamard(wires=0)
     qml.CNOT(wires=[0, 1])
@@ -382,6 +383,8 @@ class TestTransformDispatcher:  # pylint: disable=too-many-public-methods
         qfunc_transformed = dispatched_transform(qfunc_circuit, 0)
         assert callable(qfunc_transformed)
 
+        assert inspect.signature(qfunc_transformed) == inspect.signature(qfunc_circuit)
+
         with qml.tape.QuantumTape() as transformed_tape:
             qfunc_transformed(0.42)