Error when using catalyst.grad with BlockEncode

Catalyst
1

I am encountering a TracerArrayConversionError when trying to compute the gradient of a QNode that utilizes qml.BlockEncode inside a catalyst.qjit compiled function.

Example:

import pennylane as qml
import catalyst
from jax import numpy as jnp

A = jnp.array([[1, 0], [0, 1]])

dev = qml.device("lightning.qubit", wires=2)
@qml.qnode(dev)
def loss_fn(angle):
    qml.BlockEncode(A, wires=[0, 1])
    qml.RZ(angle, wires=0)
    return qml.expval(qml.Z(0))

angle = jnp.array(0.5)

grad_fn = catalyst.qjit(catalyst.grad(loss_fn))
grads = grad_fn(angle)

>>>

The above gives the error: TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on traced array with shape float64[4], which goes away if the line with BlockEncode is commented out.

Versions:

Python: 3.13.4
PennyLane version: 0.44.0
Catalyst version: 0.14.0
JAX version: 0.7.1

System:

MacOS 15.6.1
2

Thank you for the reporting this issue @kinianlo. It appears to be a bug introduced in a recent PR, we’ll get that fixed asap! I’ve open a bug report for this issue here: [BUG] QubitUnitary decomposition fails with qjit for the 2 CNOT case · Issue #9016 · PennyLaneAI/pennylane · GitHub

1 Like