abs rrule

src/UnitfulChainRules.jl

@@ -5,13 +5,17 @@ using Unitful: Quantity, Units, NoDims, FreeUnits
 using ChainRulesCore: NoTangent, @scalar_rule
 import ChainRulesCore: rrule, frule, ProjectTo
 
+const REALCOMPLEX = Union{Real, Complex}
+
 include("./rrules.jl")
 include("./frules.jl")
 include("./projection.jl")
 
 
-include("./extras.jl") # Unitful-specific rules
+include("./extras.jl") # extra Unitful-specific rules
 
 include("./trig.jl") # sin, cos, tan, etc for degrees
 
+include("./math.jl") # other math 
+
 end # module

src/math.jl

@@ -0,0 +1,8 @@
+function rrule(::typeof(abs), x::Unitful.Quantity{T,D,U}) where {T<:REALCOMPLEX, D, U}
+    Ω = abs(x)
+    function abs_pullback(ΔΩ)
+        signx = isreal(x) ? sign(x) : x / ifelse(iszero(x), oneunit(Ω), Ω)
+        return (NoTangent(), signx * real(ΔΩ))
+    end
+    return Ω, abs_pullback
+end

test/math.jl

@@ -0,0 +1,12 @@
+using Unitful
+using UnitfulChainRules
+
+using ChainRulesCore
+
+@testset "abs" begin
+    z = (1 + im)u"W"
+    Ω, pb = rrule(abs, z)
+
+    @test Ω ≈ sqrt(2)u"W"
+    @test last(pb(1.0)) ≈ (1 + im)/sqrt(2)
+end