TypeError because self._gradient_transform becomes str when using replace on execution_config in device preprocess()

[cvjjm]

Apologies for the cryptic title. I currently do not know how to describe the issue better.

I am trying to understand how to make pennylane >= 0.38 use device gradients when I stumbled across this strange behavior.

In the code below I declare that my custom device supports first order derivatives in supports_derivatives() and then want to make the necessary changes to the execution_config in preprocess(). If I manually copy the execution_config and set execution_config.use_device_gradient = True all wors as it should. But when I use replace (as is done in device_api.py) I get a very cryptic TypeError: 'str' object is not callable error deep inside the jacobian_products workflow:

from dataclasses import replace

import copy

import pennylane as qml
from pennylane.devices import DefaultQubit
from pennylane import numpy as np

from pennylane.devices.execution_config import DefaultExecutionConfig, ExecutionConfig
from pennylane.transforms.core import TransformProgram
from pennylane.tape import QuantumScript, QuantumScriptOrBatch
from numbers import Number

class MyQubit(DefaultQubit):
    def supports_derivatives(self, execution_config, circuit):
        if execution_config.gradient_method and execution_config.derivative_order == 1:
            return True

        return False

    def preprocess(self, execution_config: ExecutionConfig = DefaultExecutionConfig) -> tuple[TransformProgram, ExecutionConfig]:
        if execution_config.gradient_method and execution_config.derivative_order == 1:
            # does not work:
            replace(execution_config, use_device_gradient=True)
            # works:
            # execution_config = copy.deepcopy(execution_config)
            # execution_config.use_device_gradient = True
        return qml.transforms.core.TransformProgram(), execution_config

    def jacobian(self, *args, **kwargs):
        raise NotImplementedError("I want this to raise!")

    def compute_derivatives(
        self,
        circuits: QuantumScriptOrBatch,
        execution_config: ExecutionConfig = DefaultExecutionConfig,
    ):
        raise Exception("compute_derivatives() was called as it should! Success!")


dev = MyQubit(wires=1)

@qml.qnode(dev, diff_method='device', grad_on_execution=False)
def qnode(params):
    qml.RY(params[0], wires=[0])
    return qml.expval(qml.Hermitian(np.array([[1, 0], [0, 0]]), wires=0))

params = np.array([0.2])

print(qnode(params))
try:
    print(qml.jacobian(qnode)(params))
except Exception as exception:
    if "compute_derivatives() was called as it should! Success!" in str(exception):
        pass
    else:
        raise Exception("some other error happend") from exception
else:
    raise Exception("compute_derivatives() was not called!")

Output is:

0.9900332889206209
Traceback (most recent call last):
  File "/home/.../src/covqcstack/.../covvqetools/pl38_test_1.py", line 54, in <module>
    print(qml.jacobian(qnode)(params))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/.../src/covqcstack/.../pennylane/pennylane/_grad.py", line 456, in _jacobian_function
    jac = tuple(_jacobian(func, arg)(*args, **kwargs) for arg in _argnum)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/.../src/covqcstack/.../pennylane/pennylane/_grad.py", line 456, in <genexpr>
    jac = tuple(_jacobian(func, arg)(*args, **kwargs) for arg in _argnum)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/.../src/covqcstack/.../autograd/autograd/wrap_util.py", line 20, in nary_f
    return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/.../src/covqcstack/.../autograd/autograd/differential_operators.py", line 64, in jacobian
    return np.reshape(np.stack(grads), jacobian_shape)
                      ^^^^^^^^^^^^^^^
  File "/home/.../src/covqcstack/.../autograd/autograd/numpy/numpy_wrapper.py", line 88, in stack
    arrays = [array(arr) for arr in arrays]
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/.../src/covqcstack/.../autograd/autograd/numpy/numpy_wrapper.py", line 88, in <listcomp>
    arrays = [array(arr) for arr in arrays]
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/.../src/covqcstack/.../autograd/autograd/core.py", line 14, in vjp
    def vjp(g): return backward_pass(g, end_node)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/.../src/covqcstack/.../autograd/autograd/core.py", line 21, in backward_pass
    ingrads = node.vjp(outgrad[0])
              ^^^^^^^^^^^^^^^^^^^^
  File "/home/.../src/covqcstack/.../autograd/autograd/core.py", line 67, in <lambda>
    return lambda g: (vjp(g),)
                      ^^^^^^
  File "/home/.../src/covqcstack/.../pennylane/pennylane/workflow/interfaces/autograd.py", line 199, in grad_fn
    vjps = jpc.compute_vjp(tapes, dy)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/.../src/covqcstack/.../pennylane/pennylane/workflow/jacobian_products.py", line 297, in compute_vjp
    jacs = self.compute_jacobian(tapes)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/.../src/covqcstack/.../pennylane/pennylane/workflow/jacobian_products.py", line 327, in compute_jacobian
    jac_tapes, batch_post_processing = self._gradient_transform(tapes, **self._gradient_kwargs)
                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: 'str' object is not callable

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/.../src/covqcstack/.../covvqetools/pl38_test_1.py", line 59, in <module>
    raise Exception("some other error happend") from exception
Exception: some other error happend

This happens with both v0.38.1 and with the current master.

[CatalinaAlbornoz]

Hi @cvjjm ,
Thank you for posting this! We will investigate and get back to you. Is this blocking work?

Thinking of alternatives, you mentioned you're trying to understand how to make pennylane >= 0.38 use device gradients. Is there a specific reason why you would prefer to use replace? Or would it be ok if we found an alternative?

[cvjjm]

This is not blocking (see workaround above). I just don't understand he behavior and suspect that this is a symptom of some sort of bug that may cause unexpected behavior in other places as well and thus wanted to report it.

[dwierichs]

Hi @cvjjm,
@JerryChen97 found a solution: replace does not do in-place mutation, but returns the mutated object, so the following should be a fix (at least the MWE seems to work with it :) )

execution_config = replace(execution_config, use_device_gradient=True)

[cvjjm]

facepalm, yes that will do the trick. I was so baffled by the strange error message about self._gradient_transform being a string that I didn't consider obvious mistakes...

Out of curiosity: Why is self._gradient_transform sometimes a string and not a callable?

[dwierichs]

Out of curiosity: Why is self._gradient_transform sometimes a string and not a callable?

I might be missing a detail here, but I think it is because this attribute is passed down from the gradient method specified in a workflow. So in this case, TransformJacobianProducts._gradient_transform is "device", but the execution config (before you update it) says use_device_gradient=False, which yields an invalid configuration.

[cvjjm]

Would ne nice to check the exception config for consistency and throw a more useful error message. Just a suggestion. Closing this issue. Sorry for the noise...

[CatalinaAlbornoz]

Thanks for the feedback @cvjjm ! We are indeed considering adding some extra validation to ensure we raise a clear error in these situations. So there's no noise here, it's useful to see the checks that we can add to improve!