CRot cannot be differentiated with the adjoint method

[dime10]

Test circuit:

import numpy as np
import pennylane as qml
from catalyst import *

@qml.qnode(qml.device(backend, wires=2), diff_method="adjoint")
def f(x):
    qml.RY(0.321, wires=0)
    qml.CRot(x, 2 * x, 3 * x, wires=[0, 1])
    return qml.expval(0.5 * qml.Z(1) @ qml.X(0) - 0.4 * qml.Y(1) @ qml.H(0))


@qjit
def main(x: float):
    return qml.grad(f)(x)


result = main(0.1)
reference = main.original_function(qml.numpy.array(0.1))

assert np.allclose(result, reference)

[dime10]

The first issue one encounters here is due to the recent changes of decomposing based on differentiability, which for the special case of ops that are controlled instances whose base gate supports controlling, will be turned into a controlled version of the base gate. This can lead to infinite recursion of the result is rejected due to differentiability criteria (e.g. CRot -> Controlled(Rot) -> Controlled(Rot) -> ...). This is fixed with a simple adjustment of the condition.

Next, we can get an error of a numpy array not being C-contiguous during the JAXPR->MLIR lowering of the grad primitive:

  File "/Users/davidi/work/catalyst/frontend/catalyst/jax_primitives.py", line 517, in _grad_lowering
    attr = ir.DenseElementsAttr.get(nparray, type=const_type)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
ValueError: ndarray is not C-contiguous

I have seen this error pop up before but I'm not sure what triggers it (besides the fact that the array is F-contiguous instead of C-contiguous).
One can coerce the array to be C-contiguous to bypass this error (although it would be good to find out why it happened in the first place) with a simple nparray = np.asarray(const, order='C').

The next error we obtain is the Catalyst verification of grad claiming there is a non-differentiable op in the call graph, although the error message doesn't say which op:

catalyst failed with error code 1: Failed to run pipeline: QuantumCompilationPass
Compilation failed:
main:23:10: error: 'gradient.grad' op An operation without a valid gradient was found in code reachable from the gradient operation.
Example of operations not allowed:
 * mid circuit measurements
 * callbacks
 * ZNE mitigation.
 Try setting method="fd" to directly compute the gradient with finite difference.
    %0 = "gradient.grad"(%cst, %cst_0, %cst_1, %cst_2, %cst_3, %cst_4, %c, %cst_5, %cst_6, %cst_7, %cst_8, %cst_9, %cst_10, %cst_11, %cst_12, %cst_13, %cst_14, %cst_15, %cst_16, %cst_17, %arg0) {callee = @module_f::@f, diffArgIndices = dense<20> : tensor<1xi64>, method = "auto"} : (tensor<4x4xcomplex<f64>>, tensor<4x4xcomplex<f64>>, tensor<4x4xcomplex<f64>>, tensor<4x4xcomplex<f64>>, tensor<4x4xcomplex<f64>>, tensor<2x2xf64>, tensor<2xi64>, tensor<4x4xcomplex<f64>>, tensor<4x4xf64>, tensor<4x4xcomplex<f64>>, tensor<4x4xcomplex<f64>>, tensor<4x4xcomplex<f64>>, tensor<4x4xcomplex<f64>>, tensor<4x4xcomplex<f64>>, tensor<4x4xcomplex<f64>>, tensor<4x4xcomplex<f64>>, tensor<4x4xcomplex<f64>>, tensor<4x4xcomplex<f64>>, tensor<1xf64>, tensor<1xf64>, tensor<f64>) -> tensor<f64>

[dime10]

Update: It looks like the problem is that the decomposition of CRot generates a custom_call, and we cannot differentiate those.