Tp/fix inf to pm1 nans

Project.toml

@@ -8,6 +8,7 @@ projects = ["test", "docs"]
 
 [deps]
 ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
@@ -16,12 +17,20 @@ InverseFunctions = "3587e190-3f89-42d0-90ee-14403ec27112"
 OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
+[weakdeps]
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+
+[extensions]
+SpectralKitEnzymeExt = "Enzyme"
+
 [compat]
 ArgCheck = "1, 2"
+ChainRulesCore = "1.26.1"
 Compat = "4.18.0"
 ConcreteStructs = "0.2.7"
 ConstructionBase = "1.6.0"
 DocStringExtensions = "0.8, 0.9"
+Enzyme = "0.12, 0.13"
 InverseFunctions = "0.1"
 OrderedCollections = "1, 2"
 StaticArrays = "1"

ext/SpectralKitEnzymeExt.jl

@@ -0,0 +1,41 @@
+module SpectralKitEnzymeExt
+
+using Enzyme
+import Enzyme.EnzymeRules: forward, needs_shadow, needs_primal, augmented_primal, reverse
+import SpectralKit as SK
+
+function forward(config::EnzymeRules.FwdConfigWidth{1}, ::Const{typeof(SK._inf_to_pm1)},
+                 ::Type{RT}, z::Union{Const,Duplicated},
+                 s::Union{Const,Duplicated}) where {RT<:Union{Const,Duplicated,DuplicatedNoNeed}}
+    intermediate = SK._inf_to_pm1_intermediate(z.val, s.val)
+    if needs_shadow(config)
+        dx_dz, dx_ds = SK._inf_to_pm1_partials(intermediate)
+        tangent = zero(intermediate.x)
+        z isa Duplicated && (tangent += dx_dz * z.dval)
+        s isa Duplicated && (tangent += dx_ds * s.dval)
+        needs_primal(config) ? Duplicated(intermediate.x, tangent) : tangent
+    else
+        needs_primal(config) ? intermediate.value : nothing
+    end
+end
+
+function EnzymeRules.augmented_primal(config::EnzymeRules.RevConfigWidth{1},
+                                      ::Const{typeof(SK._inf_to_pm1)}, ::Type{RT},
+                                      z::Union{Const,Active},
+                                      s::Union{Const,Active}) where {RT<:Union{Const,Active}}
+    intermediate = SK._inf_to_pm1_intermediate(z.val, s.val)
+    primal = EnzymeRules.needs_primal(config) ? intermediate.x : nothing
+    EnzymeRules.AugmentedReturn(primal, nothing, intermediate)
+end
+
+function EnzymeRules.reverse(config::EnzymeRules.RevConfigWidth{1},
+                             ::Const{typeof(SK._inf_to_pm1)}, dret, intermediate,
+                             z::Union{Const,Active}, s::Union{Const,Active})
+    dret isa Active || return (nothing, nothing)
+    dx_dz, dx_ds = SK._inf_to_pm1_partials(intermediate)
+    dz = z isa Active ? dx_dz * dret.val : nothing
+    ds = s isa Active ? dx_ds * dret.val : nothing
+    (dz, ds)
+end
+
+end

src/SpectralKit.jl

@@ -1,6 +1,7 @@
 module SpectralKit
 
 using ArgCheck: @argcheck
+import ChainRulesCore
 using ConcreteStructs: @concrete
 using DocStringExtensions: FUNCTIONNAME, SIGNATURES, TYPEDEF
 using OrderedCollections: OrderedSet

src/transformations.jl

@@ -250,14 +250,60 @@ InfRational(; center::Real = 0.0, scale::Real = 1.0) = InfRational(promote(cente
 
 transform_from(::PM1, T::InfRational, x::Real) = T.center + T.scale * x / √(1 - abs2(x))
 
+"""
+$(SIGNATURES)
+
+Intermediate results for calculating ``x = z / √(z²+s²)``.
+
+Return a `NamedTuple` with the following values.
+
+- `m`, which is `max(abs(z), abs(s))`
+- `z_scaled`: `z/m`
+- `h_scaled`: mathematically equivalent to `hypot(z, s)/m`
+- `x`: the result.
+"""
+function _inf_to_pm1_intermediate(z::Real, s::Real)
+    z, s = promote(float(z), float(s))
+    e = exponent(max(abs(z), abs(s)))
+    z_scaled = ldexp(z, -e)
+    s_scaled = ldexp(s, -e)
+    h_scaled = hypot(z_scaled, s_scaled)
+    x = z_scaled / h_scaled
+    (; x, e, s_scaled, h_scaled)
+end
+
+function _inf_to_pm1_partials(intermediate)
+    (; x, e, s_scaled, h_scaled) = intermediate
+    s_over_h = s_scaled / h_scaled
+    d = s_over_h / h_scaled        # common term, |d| ≤ 1
+    dx_dz = ldexp(s_over_h, -e) * d # group divisions against underflow
+    dx_ds = -ldexp(x, -e) * d
+    dx_dz, dx_ds
+end
+
+_inf_to_pm1(z::Real, s::Real) = _inf_to_pm1_intermediate(z, s).x
+
+function ChainRulesCore.frule((_, ż, ṡ), ::typeof(_inf_to_pm1), z::Real, s::Real)
+    intermediate = _inf_to_pm1_intermediate(z, s)
+    dx_dz, dx_ds = _inf_to_pm1_partials(intermediate)
+    return intermediate.x, dx_dz * ż + dx_ds * ṡ
+end
+
+function ChainRulesCore.rrule(::typeof(_inf_to_pm1), z::Real, s::Real)
+    intermediate = _inf_to_pm1_intermediate(z, s)
+    dx_dz, dx_ds = _inf_to_pm1_partials(intermediate) # could be deferred into the pullback?
+    f_pullback(ȳ) = (NoTangent(), dx_dz * ȳ, dx_ds * ȳ)
+    return intermediate.x, f_pullback
+end
+
 function transform_to(::PM1, t::InfRational, y::Real)
     (; center, scale) = t
     z = y - center
-    x = z / hypot(z, scale)
     if isinf(y)
-        y > 0 ? one(x) : -one(x)
+        x = √(one(z) / one(scale)) # for type stability
+        y > 0 ? x : -x
     else
-        x
+        _inf_to_pm1(z, scale)
     end
 end
 

test/integration_enzyme.jl

@@ -6,11 +6,42 @@ using Enzyme
 
 @testset "Enzyme integration" begin
     basis = SmolyakBasis(Chebyshev(), Interior(),
-                             ntuple(_ -> BoundedLinear(1.0, 2.0), Val(2)),
+                         ntuple(_ -> BoundedLinear(1.0, 2.0), Val(2)),
                          SmolyakLevel(total = 2, each = 2))
     g(basis, θ) = linear_combination(basis, θ, (1.5, 1.5))
     d = dimension(basis)
     x = zeros(d)
     dx = zeros(d)
-    autodiff(Reverse, g, Const(basis), Duplicated(x, dx))
+    @test @inferred(autodiff(Reverse, g, Const(basis), Duplicated(x, dx))) == ((nothing, nothing),)
+end
+
+@testset "z over hypot torture test" begin
+    V = [2.0, 3.0, 1e-180, 1e-180, 0.5e-100, 1e-100, 2e-100, 1e-50, 1e100, 1e200]
+    for s in V
+        for z in V
+            for dz in V
+                for ds in V
+                    dx_z, x = Enzyme.autodiff(Enzyme.ForwardWithPrimal, SK._inf_to_pm1,
+                                              Enzyme.Duplicated(z, dz), Enzyme.Const(s))
+                    dx_s, = Enzyme.autodiff(Enzyme.Forward, SK._inf_to_pm1,
+                                            Enzyme.Const(z), Enzyme.Duplicated(s, ds))
+                    _z = BigFloat(z)
+                    _s = BigFloat(s)
+                    _r = _z / _s
+                    _q = one(_r) + _r^2
+                    _x = _z / √(_z^2 + _s^2)
+                    _dx_dz = 1 / _s / (_q * √_q)
+                    _dx_ds = -_dx_dz * _r
+                    @test x ≈ Float64(_x)
+                    @test dx_z ≈ Float64(_dx_dz * BigFloat(dz)) atol = 100 * eps() rtol = 100 * eps()
+                    @test dx_s ≈ Float64(_dx_ds * BigFloat(ds)) atol = 100 * eps() rtol = 100 * eps()
+                    (R_dx_dz, R_dx_ds), x = Enzyme.autodiff(Enzyme.ReverseWithPrimal,
+                                                            SK._inf_to_pm1, Active,
+                                                            Active(z), Active(s))
+                    @test Float64(_dx_dz) ≈ R_dx_dz
+                    @test Float64(_dx_ds) ≈ R_dx_ds
+                end
+            end
+        end
+    end
 end