using ForwardDiff, LinearAlgebra
# Fidelity between two density matrices
function fidelity(ρ::AbstractMatrix, σ::AbstractMatrix)
sqrt_ρ = sqrt(ρ)
abs2(tr(sqrt(sqrt_ρ * σ * sqrt_ρ)))
end
# Matrix representation of a Bloch vector
function matrix_representation(r)
[(1+r[1]) (r[2]-r[3]im); (r[2]+r[3]im) (1-r[1])] ./ 2
end
# Fidelity between a density matrix and a Bloch vector
function fidelity(ρ::AbstractMatrix, r::AbstractVector)
fidelity(ρ, matrix_representation(r))
end
# Gradient of the fidelity
function ∇fidelity(ρ::AbstractMatrix, r::AbstractVector)
ForwardDiff.gradient(r -> fidelity(ρ, r), r)
end
r = [0, 0, 0]
ρ = matrix_representation(r)
∇fidelity(ρ, r)

In it, I attempt to calculate the gradient of the fidelity between two quantum states, differentiating with respect to the Bloch vector of one of them. Unfortunately, I get the error

ERROR: MethodError: no method matching eigen!(::Hermitian{Complex{ForwardDiff.Dual{…}}, Matrix{Complex{…}}}; sortby::Nothing)

which I believe means that the sqrt(::AbstractMatrix) method is not differentiable by ForwardDiff.

Could someone point me to a differentiable implementation of such a method, or propose a workaround?

Obs: This 2\times2 case is only an example, and I actually need a method that works on arbitrary dimension.

Since you only need a 2x2 matrix square root, you can use the analtyical formula, which should be differentiable by ForwardDiff etc. This will be much more efficient anyway than forming a generic (albeit Hermitian) matrix and taking the square root (via eigenvalues).

When I wrote the post, I realized that I didn’t make it explicit that the 2\times2 case was only an example, and that I actually need it to work in arbitrary dimensions. I thought I had it edited to include an observation stating that, but I must have forgotten to save it. Anyway, thank you for the response, I didn’t know that StaticArrays had those optimized methods! Happy Easter!!

If you are in high dimensions (with correspondingly lots of parameters), then you probably don’t want forward-mode AD (ala ForwardDiff.jl), as in that case the cost of the gradient scales with the function cost times the number of parameters. Instead, you want reverse-mode AD (ala Zygote.jl, ReverseDiff.jl, or Enzyme.jl), as in that case the cost of the gradient scales with the function cost, independent of the number of parameters.