pennylane: [BUG] jax.jit(jax.grad()) of a circuit with shots crashes

Expected behavior

Jitting the gradient of a QNode with a device using shots, when setting the PRNGKey leads to a crash. I would expect this to work. Below there is a snippet that easily reproduces this issue on master. Do note that if you remove the jax.jit the gradient works, but this is by accident.

I think I know what is causing the bug, but the explanation is a bit involved, I will first give you a TLDR, then I will show you exactly where the crash happens, then I will reason on what is happening there.

TLDR

The problem arises because you are storing a tracer in DEfaultQubiJax._prng_key, but you are not correctly passing this prng key as an argument of the host callback in jax_jit.py:_execute. Conceptually, you should pass as an arg of the callback the prng key like you do for the parameters.

Instead, the device and therefore the _prng_key is captured in a nested series of lambdas/functions called from the callback. Therefore when the callback is executed, he encounters a tracer object for the prng key which is not substituted with concrete values and crashes.

Observing where the crash happens

As I am not very familiar with the interiors of Pennylane, and as this crash happens inside of a callback, preventing proper stack traces from being printed, I had to resort to a very primitive way of debugging. I have added several print statements in the various functions of penny lane. You can install my copy of ‘instrumented penny lane by running’

pip install git+https://github.com/PhilipVinc/pennylane@pv/debug

Using this copy, and running the snippet below, you will see the following messages printed:

INSIDE THE non_diff_wrapper CALLBACK, EXECUTING PYTHON CODE. Called with args=([array(0.34564769), array(0.45750395)], [array([1.])])
 ...
 after batch vjp, ...
  INSIDE cache_execute called with tapes=[<QuantumTape: wires=[0, 1], params=2>, <QuantumTape: wires=[0, 1], params=2>, <QuantumTape: wires=[0, 1], params=2>, <QuantumTape: wires=[0, 1], params=2>] and kwargs={}
  doing some stuff in the wrapper
   this qubit device has prng_key=Traced<ShapedArray(uint32[2])>with<DynamicJaxprTrace(level=0/1)>
   executing the circuit
     inside circuit execute for self=<DefaultQubitJax device (wires=2, shots=1000) at 0x1342ecdf0>.circuit=<QuantumTape: wires=[0, 1], params=2>
     generating samples
      - Sampling basis states for a jax qubit device with self._prng_key=Traced<ShapedArray(uint32[2])>with<DynamicJaxprTrace(level=0/1)>
      - IF THE PRNG KEY ABOVE IS A TRACER, THIS WILL CRASH!
ERROR:absl:Outside call <jax.experimental.host_callback._CallbackWrapper object at 0x134312230> threw exception Encountered an unexpected tracer. A function transformed by JAX had a side effect, allowing for a reference to an intermediate value with type uint32[2] wrapped in a DynamicJaxprTracer to escape the scope of the transformation.
JAX transformations require that functions explicitly return their outputs, and disallow saving intermediate values to global state.

The call chain at the point of the crash is the following:

• this crash is happening after the compilation. Inside of the _execute, when executing the host callback we hit the execute_fn in the callback.

  • Notice that the callback is essentially executing python code, so this code is not jitted itself (question: do you really need a callback here?), therefore if you have tracers anywhere at this point, they will lead to crashes.

• The execute_fn is a series of wrappers around QubitDevice.batch_execute, which then calls self.execute.

• In QubitDevice.execute there is a branch that, if a finite number of shots is specified, it calls

  • generate_samples and then DefaultQubitJax.sample_basis_states.

In this function we use the ._prng_key to execute some jax random functions. But as I said before, this is all being executed inside of a callback, so there should be no tracers there! Instead, as the device was captured in some lambdas, the device has a tracer as a prngKey and leads to a crash.

Possible solution

The solution is to pass the prng key as an argument to the callback. In a sense, you’d need to do something similar to cp_tape for the prng key of the device.

However, this seems complicated to do because you are not passing the device itself as an argument to those functions, but it captured inside of lambdas (I think). But maybe someone who is more familiar with pennylane @antalszava might know how to do this?

Source code

from jax.config import config
config.update("jax_enable_x64", True)

import jax
import jax.numpy as jnp
import numpy as np

import pennylane as qml

phys_qubits = 2
pars_q      = np.random.rand(3)

def minimal_circ(params, prng_key=None):
    if prng_key is not None:
        dev = qml.device("default.qubit.jax", wires=tuple(range(phys_qubits)), shots=1000, prng_key=prng_key)
    else:
        dev = qml.device("default.qubit.jax", wires=tuple(range(phys_qubits)), shots=1000)
    @qml.qnode(dev, interface="jax",diff_method="parameter-shift")
    def _measure_operator():
        qml.RY(params[0],wires=0)
        qml.RY(params[1],wires=1)
        op = qml.Hamiltonian([1.0],[qml.PauliZ(0) @ qml.PauliZ(1)])
        return qml.expval(op)
    res = _measure_operator()
    return res

grad_fun = jax.grad(minimal_circ)

jax.jit(grad_fun)(pars_q, jax.random.PRNGKey(0))

System information

>>> import pennylane as qml; qml.about()
Name: PennyLane
Version: 0.27.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author:
Author-email:
License: Apache License 2.0
Location: /Users/filippovicentini/Documents/pythonenvs/dev-pennylane/python-3.10.7/lib/python3.10/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, retworkx, scipy, semantic-version, toml
Required-by: PennyLane-Lightning

Platform info:           macOS-13.0-x86_64-i386-64bit
Python version:          3.10.7
Numpy version:           1.23.4
Scipy version:           1.9.3
Installed devices:
- default.gaussian (PennyLane-0.27.0.dev0)
- default.mixed (PennyLane-0.27.0.dev0)
- default.qubit (PennyLane-0.27.0.dev0)
- default.qubit.autograd (PennyLane-0.27.0.dev0)
- default.qubit.jax (PennyLane-0.27.0.dev0)
- default.qubit.tf (PennyLane-0.27.0.dev0)
- default.qubit.torch (PennyLane-0.27.0.dev0)
- default.qutrit (PennyLane-0.27.0.dev0)
- null.qubit (PennyLane-0.27.0.dev0)
- lightning.qubit (PennyLane-Lightning-0.26.1)

Existing GitHub issues

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