# Using QubitUnitary with qjit raises differentiability issues

**URL:** <https://discuss.pennylane.ai/t/using-qubitunitary-with-qjit-raises-differentiability-issues/7977>\
**Category:** PennyLane Help\
**Created:** [February 6, 2025, 9:36am UTC](https://discuss.pennylane.ai/t/using-qubitunitary-with-qjit-raises-differentiability-issues/7977 "2025-02-06T09:36:51Z")\
**Posts on this page:** 1\
**Showing post:** 5

<div class="post-metadata">

**Author:** ![Pia](https://avatars.discourse-cdn.com/v4/letter/p/6a8cbe/32.png) [@Pia](https://discuss.pennylane.ai/u/Pia)\
**Post date:** [February 14, 2025, 1:12pm UTC](https://discuss.pennylane.ai/t/using-qubitunitary-with-qjit-raises-differentiability-issues/7977/5 "2025-02-14T13:12:44Z")

</div>

Dear Catalina,  
good to know. I’m switching to jax now.  
Unfortunately I directly run into the next problem which I struggle to fix myself: Combining adjoint and control seems to doesn’t work under jax. Each of them individually works without problems (and with pytorch also the combination works well).

Here is the minimal example:

```auto
import matplotlib.pyplot as plt
from catalyst import qjit
from jax import numpy as jnp
import jax
import optax

device = "lightning.qubit"
shots = 100
n_wires_state = 3
lr=0.05
tsteps=1000

dev_state = qml.device(device, wires=n_wires_state, shots=shots)
unitary = jnp.array(scst.unitary_group.rvs(8))
weights = jnp.array([0.,0.,0.])

def func_circ(weights,nq):
    qml.RY(weights, wires=nq)

@jax.jit
@qml.qnode(device=dev_state, interface="jax")#diff_method="finite-diff")
def cost_circuit(weights):
    func_circ(weights[1],1)
    qml.ctrl(qml.adjoint(func_circ),control=(0))(weights[2],2)
    qml.QubitUnitary(unitary, wires=[0,1,2])
    qml.measure(wires=1)
    return [qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1))]
    
optimizer = optax.adam(lr)

@jax.jit
def costfunc(weights): 
    res = jnp.sum(jnp.array(cost_circuit(weights)), axis=1)
    vals, counts = jnp.unique(res, return_counts=True, size=1)
    res = counts[0]/len(res)
    cost = 1 - res
    return cost  

@jax.jit
def update_step_jit(i, args):
    weights, opt_state = args
    loss_val, grads = jax.value_and_grad(costfunc)(weights)
    updates, opt_state = optimizer.update(grads, opt_state)
    weights = optax.apply_updates(weights, updates)
    def print_fn():
        jax.debug.print("Step Adam pressure {i}, cost: {loss_val}", i=i, loss_val=loss_val)
    jax.lax.cond((jnp.mod(i, 50) == 0), print_fn, lambda: None)
    return (weights, opt_state)
        
@jax.jit
def optimization_jit(params, tsteps):

    opt_state = optimizer.init(params)
    args = (params, opt_state)
    (params, opt_state) = jax.lax.fori_loop(0, tsteps, update_step_jit, args)

    return params

weights = optimization_jit(weights, tsteps)
type or paste code here

```

and this is the error message I get:

> Traceback (most recent call last):  
> File “/…/minimalExample.py”, line 50, in   
> print(costfunc(weights)) # works  
> ^^^^^^^^^^^^^^^^^  
> File “/…/minimalExample.py”, line 43, in costfunc  
> res = jnp.sum(jnp.array(cost\_circuit(weights)), axis=1)  
> ^^^^^^^^^^^^^^^^^^^^^  
> File “/…/lib/python3.12/site-packages/pennylane/workflow/qnode.py”, line 905, in **call**  
> return self.\_impl\_call(\*args, \*\*kwargs)  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/…/lib/python3.12/site-packages/pennylane/workflow/qnode.py”, line 881, in \_impl\_call  
> res = qml.execute(  
> ^^^^^^^^^^^^  
> File “/…/lib/python3.12/site-packages/pennylane/workflow/execution.py”, line 227, in execute  
> tapes, post\_processing = transform\_program(tapes)  
> ^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/…/lib/python3.12/site-packages/pennylane/transforms/core/transform\_program.py”, line 580, in **call**  
> new\_tapes, fn = transform(tape, \*targs, \*\*tkwargs)  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/…/lib/python3.12/site-packages/pennylane/gradients/parameter\_shift.py”, line 765, in \_expand\_transform\_param\_shift  
> [new\_tape], postprocessing = qml.devices.preprocess.decompose(  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/…/lib/python3.12/site-packages/pennylane/transforms/core/transform\_dispatcher.py”, line 153, in **call**  
> transformed\_tapes, processing\_fn = self.\_transform(obj, \*targs, \*\*tkwargs)  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/…/lib/python3.12/site-packages/pennylane/devices/preprocess.py”, line 408, in decompose  
> if all(stopping\_condition(op) for op in tape.operations[len(prep\_op) :]):  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> File “/…/lib/python3.12/site-packages/pennylane/devices/preprocess.py”, line 408, in   
> if all(stopping\_condition(op) for op in tape.operations[len(prep\_op) :]):  
> ^^^^^^^^^^^^^^^^^^^^^^  
> File “/…/lib/python3.12/site-packages/pennylane/gradients/parameter\_shift.py”, line 746, in \_param\_shift\_stopping\_condition  
> if not op.has\_decomposition:  
> ^^^^^^^^^^^^^^^^^^^^  
> File “/…/lib/python3.12/site-packages/pennylane/ops/op\_math/controlled.py”, line 721, in has\_decomposition  
> if \_is\_single\_qubit\_special\_unitary(self.base):  
> ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^  
> jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool…  
> The error occurred while tracing the function cost\_circuit at /superscratch/psiegl/PhD/QCProject/minimalExample.py:26 for jit. This concrete value was not available in Python because it depends on the value of the argument weights.  
> See [Errors — JAX documentation](https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError)

Is there a way to allow this combination?  
Many thanks and best regards,  
Pia

---

_[View the full topic](https://discuss.pennylane.ai/t/using-qubitunitary-with-qjit-raises-differentiability-issues/7977)._