crash pennylane

grad.py

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

import pennylane as qml

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

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()
    print("res is: ", res, type(res))
    return res

grad_fun = jax.grad(minimal_circ)

# non jitted, they both work
grad_fun(pars_q)
grad_fun(pars_q, jax.random.PRNGKey(0))

# jitted, without PRNKey it works (but is useless as the result will be a constant)
jax.jit(grad_fun)(pars_q)
# jitted with PRNGKey it crashes
jax.jit(grad_fun)(pars_q, jax.random.PRNGKey(0))