[BUG] Tree-traversal + autograd gives an Arraybox error

[isaacdevlugt]

Expected behavior

I expect a QNode with mcm_method="tree-traversal" to, when differentiated, raise an error saying that tree-traversal is not supported with gradients when shots=None.

Actual behavior

TypeError: float() argument must be a string or a real number, not 'ArrayBox'

Additional information

Tree-traversal should only be differentiable with finite-diff and finite-shots. We probably just need to add some validation somewhere to fix this.

Also, can we update this warning box in the docs here? E.g.,

The tree-traversal algorithm is only supported by the DefaultQubit device, and currently does not support just-in-time (JIT) compilation. QNodes with mcm_method="tree-traversal" can only be differentiated with diff_method="finite-diff" and finite shots.

Source code

import pennylane as qml
import pennylane.numpy as pnp

dev = qml.device("default.qubit", wires=2)

@qml.qnode(dev, mcm_method="tree-traversal")
def circ(x):
    qml.RX(x, 0)
    m0 = qml.measure(0)
    return qml.expval(qml.X(0))

df = qml.grad(circ)
print(df(pnp.array(0.1)))

Tracebacks

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[41], line 17
     14     return qml.expval(qml.X(0))
     16 df = qml.grad(circ)
---> 17 print(df(pnp.array(0.1)))

File ~/Documents/pennylane/pennylane/_grad.py:224, in grad.__call__(self, *args, **kwargs)
    221     self._forward = self._fun(*args, **kwargs)
    222     return ()
--> 224 grad_value, ans = grad_fn(*args, **kwargs)  # pylint: disable=not-callable
    225 self._forward = ans
    227 return grad_value

File ~/.virtualenvs/pennylane-catalyst/lib/python3.11/site-packages/autograd/wrap_util.py:20, in unary_to_nary.<locals>.nary_operator.<locals>.nary_f(*args, **kwargs)
     18 else:
     19     x = tuple(args[i] for i in argnum)
---> 20 return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)

File ~/Documents/pennylane/pennylane/_grad.py:242, in grad._grad_with_forward(fun, x)
    236 @staticmethod
    237 @unary_to_nary
    238 def _grad_with_forward(fun, x):
    239     """This function is a replica of ``autograd.grad``, with the only
    240     difference being that it returns both the gradient *and* the forward pass
    241     value."""
--> 242     vjp, ans = _make_vjp(fun, x)  # pylint: disable=redefined-outer-name
    244     if vspace(ans).size != 1:
    245         raise TypeError(
    246             "Grad only applies to real scalar-output functions. "
    247             "Try jacobian, elementwise_grad or holomorphic_grad."
    248         )

File ~/.virtualenvs/pennylane-catalyst/lib/python3.11/site-packages/autograd/core.py:10, in make_vjp(fun, x)
      8 def make_vjp(fun, x):
      9     start_node = VJPNode.new_root()
---> 10     end_value, end_node =  trace(start_node, fun, x)
     11     if end_node is None:
     12         def vjp(g): return vspace(x).zeros()

File ~/.virtualenvs/pennylane-catalyst/lib/python3.11/site-packages/autograd/tracer.py:10, in trace(start_node, fun, x)
      8 with trace_stack.new_trace() as t:
      9     start_box = new_box(x, t, start_node)
---> 10     end_box = fun(start_box)
     11     if isbox(end_box) and end_box._trace == start_box._trace:
     12         return end_box._value, end_box._node

File ~/.virtualenvs/pennylane-catalyst/lib/python3.11/site-packages/autograd/wrap_util.py:15, in unary_to_nary.<locals>.nary_operator.<locals>.nary_f.<locals>.unary_f(x)
     13 else:
     14     subargs = subvals(args, zip(argnum, x))
---> 15 return fun(*subargs, **kwargs)

File ~/Documents/pennylane/pennylane/workflow/qnode.py:987, in QNode.__call__(self, *args, **kwargs)
    985 if qml.capture.enabled():
    986     return qml.capture.qnode_call(self, *args, **kwargs)
--> 987 return self._impl_call(*args, **kwargs)

File ~/Documents/pennylane/pennylane/workflow/qnode.py:977, in QNode._impl_call(self, *args, **kwargs)
    974     self._interface = interface
    976 try:
--> 977     res = self._execution_component(args, kwargs)
    978 finally:
    979     if old_interface == "auto":

File ~/Documents/pennylane/pennylane/workflow/qnode.py:935, in QNode._execution_component(self, args, kwargs)
    932 interface = None if self.interface == "numpy" else self.interface
    934 # pylint: disable=unexpected-keyword-arg
--> 935 res = qml.execute(
    936     (self._tape,),
    937     device=self.device,
    938     gradient_fn=gradient_fn,
    939     interface=interface,
    940     transform_program=full_transform_program,
    941     inner_transform=inner_transform_program,
    942     config=config,
    943     gradient_kwargs=gradient_kwargs,
    944     **execute_kwargs,
    945 )
    946 res = res[0]
    948 # convert result to the interface in case the qfunc has no parameters

File ~/Documents/pennylane/pennylane/workflow/execution.py:523, in execute(tapes, device, gradient_fn, interface, transform_program, inner_transform, config, grad_on_execution, gradient_kwargs, cache, cachesize, max_diff, device_vjp, mcm_config)
    521 # Exiting early if we do not need to deal with an interface boundary
    522 if no_interface_boundary_required:
--> 523     results = inner_execute(tapes)
    524     return post_processing(results)
    526 if config.use_device_jacobian_product and interface in jpc_interfaces:

File ~/Documents/pennylane/pennylane/workflow/execution.py:202, in _make_inner_execute.<locals>.inner_execute(tapes, **_)
    199 transformed_tapes, transform_post_processing = transform_program(tapes)
    201 if transformed_tapes:
--> 202     results = device.execute(transformed_tapes, execution_config=execution_config)
    203 else:
    204     results = ()

File ~/Documents/pennylane/pennylane/devices/modifiers/simulator_tracking.py:30, in _track_execute.<locals>.execute(self, circuits, execution_config)
     28 @wraps(untracked_execute)
     29 def execute(self, circuits, execution_config=DefaultExecutionConfig):
---> 30     results = untracked_execute(self, circuits, execution_config)
     31     if isinstance(circuits, QuantumScript):
     32         batch = (circuits,)

File ~/Documents/pennylane/pennylane/devices/modifiers/single_tape_support.py:32, in _make_execute.<locals>.execute(self, circuits, execution_config)
     30     is_single_circuit = True
     31     circuits = (circuits,)
---> 32 results = batch_execute(self, circuits, execution_config)
     33 return results[0] if is_single_circuit else results

File ~/Documents/pennylane/pennylane/logging/decorators.py:61, in log_string_debug_func.<locals>.wrapper_entry(*args, **kwargs)
     54     s_caller = "::L".join(
     55         [str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]]
     56     )
     57     lgr.debug(
     58         f"Calling {f_string} from {s_caller}",
     59         **_debug_log_kwargs,
     60     )
---> 61 return func(*args, **kwargs)

File ~/Documents/pennylane/pennylane/devices/default_qubit.py:642, in DefaultQubit.execute(self, circuits, execution_config)
    639 prng_keys = [self.get_prng_keys()[0] for _ in range(len(circuits))]
    641 if max_workers is None:
--> 642     return tuple(
    643         _simulate_wrapper(
    644             c,
    645             {
    646                 "rng": self._rng,
    647                 "debugger": self._debugger,
    648                 "interface": interface,
    649                 "state_cache": self._state_cache,
    650                 "prng_key": _key,
    651                 "mcm_method": execution_config.mcm_config.mcm_method,
    652                 "postselect_mode": execution_config.mcm_config.postselect_mode,
    653             },
    654         )
    655         for c, _key in zip(circuits, prng_keys)
    656     )
    658 vanilla_circuits = convert_to_numpy_parameters(circuits)[0]
    659 seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))

File ~/Documents/pennylane/pennylane/devices/default_qubit.py:643, in <genexpr>(.0)
    639 prng_keys = [self.get_prng_keys()[0] for _ in range(len(circuits))]
    641 if max_workers is None:
    642     return tuple(
--> 643         _simulate_wrapper(
    644             c,
    645             {
    646                 "rng": self._rng,
    647                 "debugger": self._debugger,
    648                 "interface": interface,
    649                 "state_cache": self._state_cache,
    650                 "prng_key": _key,
    651                 "mcm_method": execution_config.mcm_config.mcm_method,
    652                 "postselect_mode": execution_config.mcm_config.postselect_mode,
    653             },
    654         )
    655         for c, _key in zip(circuits, prng_keys)
    656     )
    658 vanilla_circuits = convert_to_numpy_parameters(circuits)[0]
    659 seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))

File ~/Documents/pennylane/pennylane/devices/default_qubit.py:908, in _simulate_wrapper(circuit, kwargs)
    907 def _simulate_wrapper(circuit, kwargs):
--> 908     return simulate(circuit, **kwargs)

File ~/Documents/pennylane/pennylane/logging/decorators.py:61, in log_string_debug_func.<locals>.wrapper_entry(*args, **kwargs)
     54     s_caller = "::L".join(
     55         [str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]]
     56     )
     57     lgr.debug(
     58         f"Calling {f_string} from {s_caller}",
     59         **_debug_log_kwargs,
     60     )
---> 61 return func(*args, **kwargs)

File ~/Documents/pennylane/pennylane/devices/qubit/simulate.py:349, in simulate(circuit, debugger, state_cache, **execution_kwargs)
    347 if has_mcm:
    348     if execution_kwargs.get("mcm_method", None) == "tree-traversal":
--> 349         return simulate_tree_mcm(circuit, prng_key=prng_key, **execution_kwargs)
    351     results = []
    352     aux_circ = qml.tape.QuantumScript(
    353         circuit.operations,
    354         circuit.measurements,
    355         shots=[1],
    356     )

File ~/Documents/pennylane/pennylane/devices/qubit/simulate.py:535, in simulate_tree_mcm(circuit, debugger, **execution_kwargs)
    531 else:
    532     shots = None
    533     skip_subtree = (
    534         stack.probs[depth] is not None
--> 535         and float(stack.probs[depth][mcm_current[depth]]) <= PROBS_TOL
    536     )
    537 # Update active branch dict
    538 invalid_postselect = (
    539     depth > 0
    540     and mcms[depth].postselect is not None
    541     and mcm_current[depth] != mcms[depth].postselect
    542 )

TypeError: float() argument must be a string or a real number, not 'ArrayBox'

System information

Name: PennyLane
Version: 0.39.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
Location: /Users/isaac/.virtualenvs/pennylane-catalyst/lib/python3.11/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, toml, typing-extensions
Required-by: PennyLane-Catalyst, PennyLane_Lightning

Platform info:           macOS-15.1-arm64-arm-64bit
Python version:          3.11.9
Numpy version:           1.26.4
Scipy version:           1.12.0
Installed devices:
- default.clifford (PennyLane-0.39.0)
- default.gaussian (PennyLane-0.39.0)
- default.mixed (PennyLane-0.39.0)
- default.qubit (PennyLane-0.39.0)
- default.qutrit (PennyLane-0.39.0)
- default.qutrit.mixed (PennyLane-0.39.0)
- default.tensor (PennyLane-0.39.0)
- null.qubit (PennyLane-0.39.0)
- reference.qubit (PennyLane-0.39.0)
- nvidia.custatevec (PennyLane-Catalyst-0.9.0)
- nvidia.cutensornet (PennyLane-Catalyst-0.9.0)
- oqc.cloud (PennyLane-Catalyst-0.9.0)
- softwareq.qpp (PennyLane-Catalyst-0.9.0)
- 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.