[BUG] [can't fix] coeff * op errors out with autograd derivatives

[Qottmann]

Expected behavior

H = X(0) + Y(0) + Z(0)

dev = qml.device("default.qubit")

@qml.qnode(dev)
def func(x, y, z):
    qml.exp(-1j * x * X(0))
    qml.exp(-1j * y * Y(0))
    qml.exp(-1j * z * Z(0))
    return qml.expval(H)

x = jnp.array(0.5)
y = jnp.array(0.5)
z = jnp.array(0.5)

jax.grad(func)(x, y, z)
# Array(-2.5406466, dtype=float32, weak_type=True)

Actual behavior

import pennylane as qml
import pennylane.numpy as pnp
from pennylane import X, Y, Z, I

H = X(0) + Y(0) + Z(0)

dev = qml.device("default.qubit")

@qml.qnode(dev)
def func(x, y, z):
    qml.exp(-1j * x * X(0))
    qml.exp(-1j * y * Y(0))
    qml.exp(-1j * z * Z(0))
    return qml.expval(H)

x = pnp.array(0.5, requires_grad=True)
y = pnp.array(0.5, requires_grad=True)
z = pnp.array(0.5, requires_grad=True)

qml.grad(func)(x, y, z)

...
Cell In[6], line 7
      5 @qml.qnode(dev)
      6 def func(x, y, z):
----> 7     qml.exp(-1j * x * X(0))
      8     qml.exp(-1j * y * Y(0))
      9     qml.exp(-1j * z * Z(0))

File ~/virtualenvs/pennylane/lib/python3.11/site-packages/autograd/numpy/numpy_boxes.py:28, in ArrayBox.__mul__(self, other)
---> 28 def __mul__(self, other): return anp.multiply(self, other)

File ~/virtualenvs/pennylane/lib/python3.11/site-packages/autograd/tracer.py:46, in primitive.<locals>.f_wrapped(*args, **kwargs)
     44     ans = f_wrapped(*argvals, **kwargs)
     45     node = node_constructor(ans, f_wrapped, argvals, kwargs, argnums, parents)
---> 46     return new_box(ans, trace, node)
     47 else:
     48     return f_raw(*args, **kwargs)

File ~/virtualenvs/pennylane/lib/python3.11/site-packages/autograd/tracer.py:120, in new_box(value, trace, node)
    118     return box_type_mappings[type(value)](value, trace, node)
    119 except KeyError:
--> 120     raise TypeError("Can't differentiate w.r.t. type {}".format(type(value)))

TypeError: Can't differentiate w.r.t. type <class 'pennylane.ops.op_math.sprod.SProd'>

Additional information

Would expect qml.exp to be differentiable

Source code

Tracebacks

System information

Name: PennyLane
Version: 0.41.0.dev18
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/korbinian.kottmann/virtualenvs/pennylane/lib/python3.11/site-packages
Editable project location: /Users/korbinian.kottmann/Xanadu/pennylane
Requires: appdirs, autograd, autoray, cachetools, diastatic-malt, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, tomlkit, typing_extensions
Required-by: PennyLane_Lightning

Platform info:           macOS-15.3.1-arm64-arm-64bit
Python version:          3.11.11
Numpy version:           1.26.4
Scipy version:           1.15.1
Installed devices:
- default.clifford (PennyLane-0.41.0.dev18)
- default.gaussian (PennyLane-0.41.0.dev18)
- default.mixed (PennyLane-0.41.0.dev18)
- default.qubit (PennyLane-0.41.0.dev18)
- default.qutrit (PennyLane-0.41.0.dev18)
- default.qutrit.mixed (PennyLane-0.41.0.dev18)
- default.tensor (PennyLane-0.41.0.dev18)
- null.qubit (PennyLane-0.41.0.dev18)
- reference.qubit (PennyLane-0.41.0.dev18)
- lightning.qubit (PennyLane_Lightning-0.40.0)

Existing GitHub issues

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

[albi3ro]

Note that this also occurs with things like:

@qml.qnode(dev)
def circuit(coeff):
    return qml.expval(coeff * qml.Z(0))

qml.grad(circuit)(qml.numpy.array(2.0))

But you can just do qml.s_prod(-1j * x, X(0)) and get it to work.

Given this issue is in the dunder methods for a ArrayBox, I'm not sure it's possible for us to fix without patching autograd.

[Qottmann]

Ah okay I didnt realize this was an autograd issue!

But you can just do qml.s_prod(-1j * x, X(0)) and get it to work.

That will do, thanks!

[albi3ro]

Potentially worth keeping open as a "cant fix", and renaming it to something like "coeff * op errors out with autograd derivatives".