From 4acca58835812f449702bc5726415fb01fd9dd73 Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Thu, 4 Sep 2025 09:48:24 -0400 Subject: [PATCH 1/7] Add ChainRulesCore extension --- Project.toml | 7 +++++++ ext/VectorInterfaceChainRulesCoreExt.jl | 3 +++ 2 files changed, 10 insertions(+) create mode 100644 ext/VectorInterfaceChainRulesCoreExt.jl diff --git a/Project.toml b/Project.toml index c1b249a..8461405 100644 --- a/Project.toml +++ b/Project.toml @@ -6,8 +6,15 @@ version = "0.5" [deps] LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +[weakdeps] +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + +[extensions] +VectorInterfaceChainRulesCoreExt = "ChainRulesCore" + [compat] Aqua = "0.6, 0.7, 0.8" +ChainRulesCore = "1" LinearAlgebra = "1" Test = "1" TestExtras = "0.2,0.3" diff --git a/ext/VectorInterfaceChainRulesCoreExt.jl b/ext/VectorInterfaceChainRulesCoreExt.jl new file mode 100644 index 0000000..2d52162 --- /dev/null +++ b/ext/VectorInterfaceChainRulesCoreExt.jl @@ -0,0 +1,3 @@ +module VectorInterfaceChainRulesCoreExt + +end From 30ca7dcc33b6a9fbcb592f59125e56620f82d737 Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Thu, 4 Sep 2025 09:49:21 -0400 Subject: [PATCH 2/7] Add rrules --- ext/VectorInterfaceChainRulesCoreExt.jl | 65 +++++++++++++++++++++++++ 1 file changed, 65 insertions(+) diff --git a/ext/VectorInterfaceChainRulesCoreExt.jl b/ext/VectorInterfaceChainRulesCoreExt.jl index 2d52162..cc9c927 100644 --- a/ext/VectorInterfaceChainRulesCoreExt.jl +++ b/ext/VectorInterfaceChainRulesCoreExt.jl @@ -1,3 +1,68 @@ module VectorInterfaceChainRulesCoreExt +using VectorInterface +using ChainRulesCore: ChainRulesCore, NoTangent, ZeroTangent, StructuralTangent, rrule, + backing, unthunk + +# scale +# ----- +function ChainRulesCore.rrule(::typeof(scale), x, α::Number) + function scale_pullback(Δx_) + Δx = unthunk(Δx_) + return NoTangent(), scale(Δx, conj(α)), inner(x, Δx) + end + return scale(x, α), scale_pullback +end + +ChainRulesCore.rrule(::typeof(scale!!), x, α::Number) = rrule(scale, x, α) + +function ChainRulesCore.rrule(::typeof(scale!!), y, x, α::Number) + function scale_pullback(Δy_) + Δy = unthunk(Δy_) + return NoTangent(), ZeroTangent(), scale(Δy, conj(α)), inner(x, Δy) + end + return scale!!(y, x, α), scale_pullback +end + +# add +# --- +function ChainRulesCore.rrule(::typeof(add), y, x, α::Number, β::Number) + z = add(y, x, α, β) + function add_pullback(Δz_) + Δz = unthunk(Δz_) + return NoTangent(), + scale(Δz, conj(β)), scale(Δz, conj(α)), + inner(x, Δz), inner(y, Δz) + end + return z, add_pullback +end + +ChainRulesCore.rrule(::typeof(add!!), y, x, α::Number, β::Number) = rrule(add, y, x, α, β) + +# inner +# ----- +function ChainRulesCore.rrule(::typeof(inner), x, y) + function inner_pullback(Δn_) + Δn = unthunk(Δn_) + return NoTangent(), scale(y, conj(Δn)), scale(x, Δn) + end + return inner(x, y), inner_pullback +end + +# Tangent support +# --------------- +VectorInterface.scale(x::StructuralTangent, α::Number) = map(Base.Fix2(scale, α), x) + +function VectorInterface.add(y::StructuralTangent{P}, x::StructuralTangent{P}, α::Number, + β::Number) where {P} + return ChainRulesCore.add!!(scale(y, β), scale(x, α)) +end + +function VectorInterface.inner(y::StructuralTangent{P}, x::P) where {P} + return inner(backing(y), backing(x)) +end +function VectorInterface.inner(y::P, x::StructuralTangent{P}) where {P} + return inner(backing(y), backing(x)) +end + end From 57397a94042f569db9eabc3c45109cb0cda91b31 Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Thu, 4 Sep 2025 09:49:24 -0400 Subject: [PATCH 3/7] Add tests --- Project.toml | 4 +- test/chainrules.jl | 95 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 98 insertions(+), 1 deletion(-) create mode 100644 test/chainrules.jl diff --git a/Project.toml b/Project.toml index 8461405..24265bd 100644 --- a/Project.toml +++ b/Project.toml @@ -15,6 +15,7 @@ VectorInterfaceChainRulesCoreExt = "ChainRulesCore" [compat] Aqua = "0.6, 0.7, 0.8" ChainRulesCore = "1" +ChainRulesTestUtils = "1" LinearAlgebra = "1" Test = "1" TestExtras = "0.2,0.3" @@ -22,8 +23,9 @@ julia = "1" [extras] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a" [targets] -test = ["Test", "TestExtras", "Aqua"] +test = ["Test", "TestExtras", "Aqua", "ChainRulesTestUtils", "ChainRulesCore"] diff --git a/test/chainrules.jl b/test/chainrules.jl new file mode 100644 index 0000000..9433ecf --- /dev/null +++ b/test/chainrules.jl @@ -0,0 +1,95 @@ +module ChainRules + +using VectorInterface +using VectorInterface: MinimalMVec, MinimalSVec, MinimalVec +using Test, TestExtras +using ChainRulesTestUtils +using ChainRulesCore: ChainRulesCore, AbstractZero + +precision(::Type{T}) where {T<:Union{Float32,ComplexF32}} = sqrt(eps(Float32)) +precision(::Type{T}) where {T<:Union{Float64,ComplexF64}} = sqrt(eps(Float64)) + +# Small adaptations to make tests work with MinimalVec +function ChainRulesTestUtils.test_approx(::AbstractZero, x::MinimalVec, msg=""; kwargs...) + return test_approx(zerovector(x), x, msg; kwargs...) +end +function ChainRulesTestUtils.test_approx(x::MinimalVec, ::AbstractZero, msg=""; kwargs...) + return test_approx(x, zerovector(x), msg; kwargs...) +end +Base.collect(x::MinimalVec) = x.vec + +eltypes = (Float32, Float64, ComplexF64) + +@testset "scale pullbacks ($T)" for T in eltypes + n = 12 + atol = rtol = n * precision(T) + + # Vector + x = randn(T, n) + y = randn(T, n) + α = randn(T) + test_rrule(scale, x, α; atol, rtol) + test_rrule(scale!!, x, α; atol, rtol) + test_rrule(scale!!, y, x, α; atol, rtol) + + # MinimalMVec + mx = MinimalMVec(x) + my = MinimalMVec(y) + test_rrule(scale, mx, α; atol, rtol, check_inferred=false) + test_rrule(scale!!, mx, α; atol, rtol, check_inferred=false) + test_rrule(scale!!, my, mx, α; atol, rtol, check_inferred=false) + + # MinimalSVec + mx = MinimalSVec(x) + my = MinimalSVec(y) + test_rrule(scale, mx, α; atol, rtol, check_inferred=false) + test_rrule(scale!!, mx, α; atol, rtol, check_inferred=false) + test_rrule(scale!!, my, mx, α; atol, rtol, check_inferred=false) +end + +@testset "add pullbacks ($T)" for T in eltypes + n = 12 + atol = rtol = n * precision(T) + + # Vector + x = randn(T, n) + y = randn(T, n) + α = randn(T) + β = randn(T) + test_rrule(add, y, x, α, β; atol, rtol) + test_rrule(add!!, y, x, α, β; atol, rtol) + + # MinimalMVec + mx = MinimalMVec(x) + my = MinimalMVec(y) + test_rrule(add, my, mx, α, β; atol, rtol, check_inferred=false) + test_rrule(add!!, my, mx, α, β; atol, rtol, check_inferred=false) + + # MinimalSVec + mx = MinimalSVec(x) + my = MinimalSVec(y) + test_rrule(add, my, mx, α, β; atol, rtol, check_inferred=false) + test_rrule(add!!, my, mx, α, β; atol, rtol, check_inferred=false) +end + +@testset "inner pullbacks ($T)" for T in eltypes + n = 12 + atol = rtol = n * precision(T) + + # Vector + x = randn(T, n) + y = randn(T, n) + test_rrule(inner, x, y; atol, rtol) + + # MinimalMVec + mx = MinimalMVec(x) + my = MinimalMVec(y) + test_rrule(inner, my, mx; atol, rtol, check_inferred=false) + + # MinimalSVec + mx = MinimalSVec(x) + my = MinimalSVec(y) + test_rrule(inner, my, mx; atol, rtol, check_inferred=false) +end + +end From 75e891e461875d95915982b5f6a3679c57958dfe Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Thu, 4 Sep 2025 09:49:40 -0400 Subject: [PATCH 4/7] Bump v0.5.1 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 24265bd..09d8ac5 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "VectorInterface" uuid = "409d34a3-91d5-4945-b6ec-7529ddf182d8" authors = ["Jutho Haegeman and contributors"] -version = "0.5" +version = "0.5.1" [deps] LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" From 1f85627b332fb278d49f556dd4cb39d545d51676 Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Thu, 4 Sep 2025 09:57:36 -0400 Subject: [PATCH 5/7] ensure backwards compatibility --- Project.toml | 1 + test/runtests.jl | 6 ++++++ 2 files changed, 7 insertions(+) diff --git a/Project.toml b/Project.toml index 09d8ac5..cc279e2 100644 --- a/Project.toml +++ b/Project.toml @@ -23,6 +23,7 @@ julia = "1" [extras] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" TestExtras = "5ed8adda-3752-4e41-b88a-e8b09835ee3a" diff --git a/test/runtests.jl b/test/runtests.jl index 76d08ad..0777e59 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -34,3 +34,9 @@ using VectorInterface using Aqua Aqua.test_all(VectorInterface) end + +@static if isdefined(Base, :get_extension) + println("Testing AD rules") + println("================") + include("chainrules.jl") +end From 81a6aee52c31ce76c87629af71d7375c2bb192db Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Thu, 4 Sep 2025 10:07:43 -0400 Subject: [PATCH 6/7] ensure forwards compatibility --- test/runtests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/runtests.jl b/test/runtests.jl index 0777e59..e99badd 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -35,7 +35,7 @@ using Aqua Aqua.test_all(VectorInterface) end -@static if isdefined(Base, :get_extension) +@static if isdefined(Base, :get_extension) && isempty(VERSION.prerelease) println("Testing AD rules") println("================") include("chainrules.jl") From 1642c5e531e1b0425c7b0c21dfe966b3380e7a3b Mon Sep 17 00:00:00 2001 From: Lukas Devos Date: Thu, 18 Sep 2025 10:59:05 +0200 Subject: [PATCH 7/7] Apply suggestions from code review Co-authored-by: Jutho --- test/chainrules.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/chainrules.jl b/test/chainrules.jl index 9433ecf..c61b04a 100644 --- a/test/chainrules.jl +++ b/test/chainrules.jl @@ -84,12 +84,12 @@ end # MinimalMVec mx = MinimalMVec(x) my = MinimalMVec(y) - test_rrule(inner, my, mx; atol, rtol, check_inferred=false) + test_rrule(inner, mx, my; atol, rtol, check_inferred=false) # MinimalSVec mx = MinimalSVec(x) my = MinimalSVec(y) - test_rrule(inner, my, mx; atol, rtol, check_inferred=false) + test_rrule(inner, mx, my; atol, rtol, check_inferred=false) end end