From 80b910813e7c65fe52f0a7e709570e85165f257f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tam=C3=A1s=20K=2E=20Papp?= Date: Mon, 19 Oct 2026 15:58:21 +0200 Subject: [PATCH 1/2] =?UTF-8?q?implement=20custom=20derivatives=20for=20z/?= =?UTF-8?q?=E2=88=9A(z=C2=B2+s=C2=B2)=20transformation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Project.toml | 9 +++++++ ext/SpectralKitEnzymeExt.jl | 41 +++++++++++++++++++++++++++++++ src/SpectralKit.jl | 1 + src/transformations.jl | 48 ++++++++++++++++++++++++++++++++++++- test/integration_enzyme.jl | 35 +++++++++++++++++++++++++-- 5 files changed, 131 insertions(+), 3 deletions(-) create mode 100644 ext/SpectralKitEnzymeExt.jl diff --git a/Project.toml b/Project.toml index 12d2a7d..f6fc1ed 100644 --- a/Project.toml +++ b/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" diff --git a/ext/SpectralKitEnzymeExt.jl b/ext/SpectralKitEnzymeExt.jl new file mode 100644 index 0000000..43dca1a --- /dev/null +++ b/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 diff --git a/src/SpectralKit.jl b/src/SpectralKit.jl index ea2198d..1869176 100644 --- a/src/SpectralKit.jl +++ b/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 diff --git a/src/transformations.jl b/src/transformations.jl index 52d1863..2e1531e 100644 --- a/src/transformations.jl +++ b/src/transformations.jl @@ -250,10 +250,56 @@ 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)) + m = max(abs(z), abs(s)) + z_scaled = z / m + s_scaled = s / m + h_scaled = hypot(z_scaled, s_scaled) + x = z_scaled / h_scaled + (; x, m, s_scaled, h_scaled) +end + +function _inf_to_pm1_partials(intermediate) + (; x, m, s_scaled, h_scaled) = intermediate + s_over_h = s_scaled / h_scaled + d = s_over_h / h_scaled # common term, |d| ≤ 1 + dx_dz = (s_over_h / m) * d # group divisions against underflow + dx_ds = -(x / m) * 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) + x = _inf_to_pm1(z, scale) if isinf(y) y > 0 ? one(x) : -one(x) else diff --git a/test/integration_enzyme.jl b/test/integration_enzyme.jl index 40bc82e..507526f 100644 --- a/test/integration_enzyme.jl +++ b/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 From 057a3d108e9827e6eb81df197a81a1aa6f2d5622 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tam=C3=A1s=20K=2E=20Papp?= Date: Mon, 19 Oct 2026 16:16:06 +0200 Subject: [PATCH 2/2] use ldexp --- src/transformations.jl | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/src/transformations.jl b/src/transformations.jl index 2e1531e..3b2ddfa 100644 --- a/src/transformations.jl +++ b/src/transformations.jl @@ -264,20 +264,20 @@ Return a `NamedTuple` with the following values. """ function _inf_to_pm1_intermediate(z::Real, s::Real) z, s = promote(float(z), float(s)) - m = max(abs(z), abs(s)) - z_scaled = z / m - s_scaled = s / m + 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, m, s_scaled, h_scaled) + (; x, e, s_scaled, h_scaled) end function _inf_to_pm1_partials(intermediate) - (; x, m, s_scaled, h_scaled) = 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 = (s_over_h / m) * d # group divisions against underflow - dx_ds = -(x / m) * d + dx_dz = ldexp(s_over_h, -e) * d # group divisions against underflow + dx_ds = -ldexp(x, -e) * d dx_dz, dx_ds end @@ -299,11 +299,11 @@ end function transform_to(::PM1, t::InfRational, y::Real) (; center, scale) = t z = y - center - x = _inf_to_pm1(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