# Is there a differentiable implementation of matrix square root?
**URL:** https://discourse.julialang.org/t/is-there-a-differentiable-implementation-of-matrix-square-root/112334
**Category:** General Usage
**Tags:** differentiation, linearalgebra, forwarddiff
**Created:** [March 30, 2024, 6:10pm UTC](https://discourse.julialang.org/t/is-there-a-differentiable-implementation-of-matrix-square-root/112334 "2024-03-30T18:10:49Z")
**Posts on this page:** 6
**Page:** 1
### Author:  [@marcsgil](https://discourse.julialang.org/u/marcsgil)
#### Post date: [March 30, 2024, 6:10pm UTC](https://discourse.julialang.org/t/is-there-a-differentiable-implementation-of-matrix-square-root/112334/1 "2024-03-30T18:10:50Z")
Hello!
I’m trying to make the following code work:
```julia
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](https://en.wikipedia.org/wiki/Fidelity_of_quantum_states) between two quantum states, differentiating with respect to the [Bloch vector](https://en.wikipedia.org/wiki/Bloch_sphere#Definition) of one of them. Unfortunately, I get the error
```julia
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`](https://github.com/JuliaDiff/ForwardDiff.jl).
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.
Thanks!
---
### Author:  [@gdalle](https://discourse.julialang.org/u/gdalle)
#### Post date: [March 31, 2024, 8:24am UTC](https://discourse.julialang.org/t/is-there-a-differentiable-implementation-of-matrix-square-root/112334/2 "2024-03-31T08:24:12Z")
You can take a look at Enzyme.jl or DifferentiableFactorizations.jl, although I’m unsure how either behaves with complex input
---
### Author:  [@stevengj](https://discourse.julialang.org/u/stevengj)
#### Post date: [March 31, 2024, 2:30pm UTC](https://discourse.julialang.org/t/is-there-a-differentiable-implementation-of-matrix-square-root/112334/3 "2024-03-31T14:30:11Z")
> [@marcsgil](#):
>
> or propose a workaround?
Since you only need a 2x2 matrix square root, you can use the [analtyical formula](https://en.wikipedia.org/wiki/Square_root_of_a_2_by_2_matrix), 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).
Actually, this 2x2 formula is [already included in StaticArrays.jl](https://github.com/JuliaArrays/StaticArrays.jl/blob/07c12450d1b3481dda4b503564ae4a5cb4e27ce4/src/sqrtm.jl#L13-L22), so you can just use an `SMatrix` — which you [should probably be using anyway](https://docs.julialang.org/en/v1/manual/performance-tips/#Consider-StaticArrays.jl-for-small-fixed-size-vector/matrix-operations) for such small fixed-size matrices — and it should work (and be _much_ faster).
Changing your code to:
```julia
using StaticArrays
function matrix_representation(r)
@SMatrix[(1+r[1]) (r[2]-r[3]im); (r[2]+r[3]im) (1-r[1])] ./ 2
end
```
gives
```julia
julia> ∇fidelity(ρ, [0,0,0])
3-element Vector{Float64}:
0.0
0.0
0.0
julia> ∇fidelity(ρ, [0.1,0.2,0.3])
3-element Vector{Float64}:
-0.053916386601719206
-0.10783277320343838
-0.1617491598051576
```
which matches a finite-difference check:
```julia
julia> r = [0.1,0.2,0.3]; dr = randn(3) * 1e-8;
julia> isapprox(fidelity(ρ,r+dr) - fidelity(ρ,r), ∇fidelity(ρ,r) ⋅ dr, rtol=1e-5)
true
```
🌸🌸🌸🐇🥚🧺🐇🥚🧺🐇🥚🧺🐇🥚🧺🌸🌸🌸
---
### Author:  [@marcsgil](https://discourse.julialang.org/u/marcsgil)
#### Post date: [March 31, 2024, 2:58pm UTC](https://discourse.julialang.org/t/is-there-a-differentiable-implementation-of-matrix-square-root/112334/4 "2024-03-31T14:58:03Z")
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!! 🐇 🐇 🐇 🥚 🥚 🥚
---
### Author:  [@stevengj](https://discourse.julialang.org/u/stevengj)
#### Post date: [March 31, 2024, 4:48pm UTC](https://discourse.julialang.org/t/is-there-a-differentiable-implementation-of-matrix-square-root/112334/5 "2024-03-31T16:48:38Z")
> [@marcsgil](#):
>
> and that I actually need it to work in arbitrary dimensions.
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.
ChainRules.jl (used by Zygote.jl) already has a rule for [differentiating the `sqrt` of a `Hermitian` matrix](https://github.com/JuliaDiff/ChainRules.jl/blob/9f1817a22404259113e230bef149a54d379a660b/src/rulesets/LinearAlgebra/symmetric.jl#L306-L313). (Don’t forget to wrap your matrix in `Hermitian`.)
---
### Author:  [@wsmoses](https://discourse.julialang.org/u/wsmoses)
#### Post date: [March 31, 2024, 10:15pm UTC](https://discourse.julialang.org/t/is-there-a-differentiable-implementation-of-matrix-square-root/112334/6 "2024-03-31T22:15:47Z")
See here for a discussion of Complex numbers in Enzyme: [FAQ · Enzyme.jl](https://enzymead.github.io/Enzyme.jl/dev/faq/#Complex-numbers)