# Enzyme Reverse Diff rules for complex sqrt
**URL:** https://discourse.julialang.org/t/enzyme-reverse-diff-rules-for-complex-sqrt/111489
**Category:** Modelling & Simulations
**Tags:** autodiff, enzyme
**Created:** [March 11, 2024, 9:14pm UTC](https://discourse.julialang.org/t/enzyme-reverse-diff-rules-for-complex-sqrt/111489 "2024-03-11T21:14:24Z")
**Posts on this page:** 7
**Page:** 1
### Author:  [@dchang10](https://discourse.julialang.org/u/dchang10)
#### Post date: [March 11, 2024, 9:14pm UTC](https://discourse.julialang.org/t/enzyme-reverse-diff-rules-for-complex-sqrt/111489/1 "2024-03-11T21:14:24Z")
I have a function that I’m trying to reverse mode diff through that at some point needs to calculate the square root of a complex number. I have been able to define custom rules for the `Forward` mode `autodiff` in the following way:
```julia
function forward(func::Const{typeof(sqrt)}, ::Type{<:Duplicated}, x::Duplicated{Complex{T}}) where {T<:Real}
ret = func.val(x.val)
return Duplicated(ret, 1/(2ret) * x.dval)
end
```
but have been unable to do the same for the `Reverse` mode. [Following the custom rules tutorial](https://enzyme.mit.edu/julia/stable/generated/custom_rule/), I have defined the `augmented_primal` and `reverse` functions like:
```julia
function augmented_primal(config::ConfigWidth{1}, func::Const{typeof(sqrt)}, ::Type{<:Active}, x::Duplicated{Complex{T}}) where {T<:Real}
println("In custom augmented primal rule.")
if needs_primal(config)
primal = func.val(x.val)
else
primal = nothing
end|
# Save x in tape if x will be overwritten
if overwritten(config)[3]
tape = copy(x.val)
else
tape = nothing
end
# Return an AugmentedReturn object with shadow = nothing
return AugmentedReturn(primal, nothing, tape)
end
function reverse(config::ConfigWidth{1}, func::Const{typeof(sqrt)}, dret::Active, tape, x::Duplicated{Complex{T}}) where {T<:Real}
println("In custom reverse rule.")
# retrieve x value, either from original x or from tape if x may have been overwritten.
xval = overwritten(config)[3] ? tape : x.val
x.dval += inv(2func(xval)) * dret.val
return (nothing, nothing)
end
```
I however receive the error :
`ERROR: Duplicated Returns not yet handled`
when trying to execute the following test function
```julia
function test(η)
ans = sqrt(η + 0im)
return abs(ans)
end
autodiff(Enzyme.Reverse, test, Duplicated, Duplicated(2.0,1.0))
```
---
### Author:  [@wsmoses](https://discourse.julialang.org/u/wsmoses)
#### Post date: [March 13, 2024, 7:35am UTC](https://discourse.julialang.org/t/enzyme-reverse-diff-rules-for-complex-sqrt/111489/2 "2024-03-13T07:35:34Z")
You probably want an active return here.
---
### Author:  [@dchang10](https://discourse.julialang.org/u/dchang10)
#### Post date: [March 15, 2024, 12:56am UTC](https://discourse.julialang.org/t/enzyme-reverse-diff-rules-for-complex-sqrt/111489/4 "2024-03-15T00:56:27Z")
I tried changing the signature to
```julia
autodiff(Enzyme.Reverse, test, Active, Duplicated(2.0,1.0))
```
But enzyme still seemed to differentiate `sqrt` without using my rule. I assume I’ve set up the method signatures incorrectly somehow
---
### Author:  [@wsmoses](https://discourse.julialang.org/u/wsmoses)
#### Post date: [March 17, 2024, 12:16am UTC](https://discourse.julialang.org/t/enzyme-reverse-diff-rules-for-complex-sqrt/111489/5 "2024-03-17T00:16:26Z")
Reverse mode requires floats to be passed in via active not duplicated
---
### Author:  [@dchang10](https://discourse.julialang.org/u/dchang10)
#### Post date: [March 17, 2024, 1:50am UTC](https://discourse.julialang.org/t/enzyme-reverse-diff-rules-for-complex-sqrt/111489/6 "2024-03-17T01:50:29Z")
Thanks. Apologies for my ignorance. I’ve tried defining the rules this way:
```julia
function augmented_primal(config::ConfigWidth{1}, func::Const{typeof(sqrt)}, ::Type{<:Active}, x::Active{Complex{T}}) where {T<:Real}
println("In custom augmented primal rule.")
if needs_primal(config)
primal = func.val(x.val)
else
primal = nothing
end
# Save x in tape if x will be overwritten
if overwritten(config)[3]
tape = copy(x.val)
else
tape = nothing
end
# Return an AugmentedReturn object with shadow = nothing
return AugmentedReturn(primal, nothing, tape)
end
function reverse(config::ConfigWidth{1}, func::Const{typeof(sqrt)}, dret::Active, tape, x::Active{Complex{T}}) where {T<:Real}
println("In custom reverse rule.")
# retrieve x value, either from original x or from tape if x may have been overwritten.
xval = overwritten(config)[3] ? tape : x.val
x.dval += inv(2func(xval)) * dret.val
return (nothing, )
end
```
but enzyme is still ignoring the square root rule on evaluation
```julia
function test(η)
ans = sqrt(η + 0im)
return abs(ans)
end
autodiff(Enzyme.Reverse, test, Active, Active(2.0))
Stacktrace:
[1] |
@ ./int.jl:372
[2] ldexp
@ ./math.jl:964
[3] sqrt
@ ./complex.jl:541
[4] test
@ ~/Software/Krang.jl/examples/mwe.jl:7
Stacktrace:
[1] throwerr(cstr::Cstring)
@ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:1289
[2] |
@ ./int.jl:372 [inlined]
[3] ldexp
@ ./math.jl:964 [inlined]
[4] sqrt
@ ./complex.jl:541 [inlined]
[5] test
@ ~/Software/Krang.jl/examples/mwe.jl:7 [inlined]
[6] diffejulia_test_4396wrap
@ ~/Software/Krang.jl/examples/mwe.jl:0
[7] macro expansion
@ ~/.julia/dev/Enzyme/src/compiler.jl:5440 [inlined]
[8] enzyme_call(::Val{…}, ::Ptr{…}, ::Type{…}, ::Type{…}, ::Val{…}, ::Type{…}, ::Type{…}, ::Const{…}, ::Type{…}, ::Active{…}, ::Float64)
@ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:5118
[9] (::Enzyme.Compiler.CombinedAdjointThunk{…})(::Const{…}, ::Active{…}, ::Vararg{…})
@ Enzyme.Compiler ~/.julia/dev/Enzyme/src/compiler.jl:5000
[10] autodiff
@ ~/.julia/dev/Enzyme/src/Enzyme.jl:0 [inlined]
[11] autodiff(mode::ReverseMode{false, FFIABI, false}, f::typeof(test), ::Type{Active}, args::Active{Float64})
@ Enzyme ~/.julia/dev/Enzyme/src/Enzyme.jl:287
```
---
### Author:  [@dchang10](https://discourse.julialang.org/u/dchang10)
#### Post date: [March 18, 2024, 9:47pm UTC](https://discourse.julialang.org/t/enzyme-reverse-diff-rules-for-complex-sqrt/111489/7 "2024-03-18T21:47:14Z")
Apparently the error was because I was not sub-typing on `Complex`. Here is a solution which works
```julia
function augmented_primal(config::ConfigWidth{1}, func::Const{typeof(sqrt)}, ::Type{<:Active}, x::Active{<:Complex{T}}) where {T<:Real}
println("In custom augmented primal rule.")
if needs_primal(config)
primal = func.val(x.val)
else
primal = nothing
end
# Save x in tape if x will be overwritten
if overwritten(config)[2]
tape = copy(x.val)
else
tape = nothing
end
# Return an AugmentedReturn object with shadow = nothing
return AugmentedReturn(primal, nothing, tape)
end
function reverse(config::ConfigWidth{1}, ::Const{typeof(sqrt)}, dret::Active, tape, x::Active{<:Complex{T}}) where {T<:Real}
println("In custom reverse rule.")
# retrieve x value, either from original x or from tape if x may have been overwritten.
xval = overwritten(config)[2] ? tape : x.val
dx = inv(2*sqrt(xval))' * dret.val
return (dx, )
end
```
```julia
function test(η)
ans = sqrt(η*exp((π/4)*1im))
return abs(ans)
end
autodiff(Enzyme.Reverse, test, Active, Active(2.0 + 0im))
```
with output :
```julia
In custom augmented primal rule.
In custom reverse rule.
((0.35355339059327373 - 1.1496735851465466e-17im,),)
```
---
### Author:  [@wsmoses](https://discourse.julialang.org/u/wsmoses)
#### Post date: [April 6, 2024, 4:32pm UTC](https://discourse.julialang.org/t/enzyme-reverse-diff-rules-for-complex-sqrt/111489/8 "2024-04-06T16:32:36Z")
Separately we’ve just added a first class complex sqrt rule on the main branch so you shouldn’t need a custom rule for this case any more