Custom differentiable UFL operator

Hello, this is a question about UFL and composing differentiable operations.
Let’s say I have a dolfinx.fem.Function p representing some parameter field, and a scalar valued UFL expression J(p). It is well known that I can evaluate the gradient of J w.r.t. p with something like

dJdp = ufl.derivative(J, p)
dJdp_vector = dolfinx.fem.assemble_vector(dolfinx.fem.form(dJdp))

Now, I wish to add some custom operation on p inside the expression of J(p). Let’s name q that operator: J(p) = J(q(p)). However, q is not a UFL expression, it is implemented with an external python code. Though I am able to compute the jacobian matrix of q for a given input p.

My question is: is it possible to wrap q inside a UFL object such that it is still possible to define a UFL expression for J, and evaluate its gradient with automatic differentiation?

You would need to come up with a minimal example of what p is. You can use: ufl package — Unified Form Language (UFL) 2021.1.0 documentation