diff --git a/Project.toml b/Project.toml index c1b249a..cc279e2 100644 --- a/Project.toml +++ b/Project.toml @@ -1,13 +1,21 @@ 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" +[weakdeps] +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" + +[extensions] +VectorInterfaceChainRulesCoreExt = "ChainRulesCore" + [compat] Aqua = "0.6, 0.7, 0.8" +ChainRulesCore = "1" +ChainRulesTestUtils = "1" LinearAlgebra = "1" Test = "1" TestExtras = "0.2,0.3" @@ -15,8 +23,10 @@ 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" [targets] -test = ["Test", "TestExtras", "Aqua"] +test = ["Test", "TestExtras", "Aqua", "ChainRulesTestUtils", "ChainRulesCore"] diff --git a/ext/VectorInterfaceChainRulesCoreExt.jl b/ext/VectorInterfaceChainRulesCoreExt.jl new file mode 100644 index 0000000..cc9c927 --- /dev/null +++ b/ext/VectorInterfaceChainRulesCoreExt.jl @@ -0,0 +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 diff --git a/test/chainrules.jl b/test/chainrules.jl new file mode 100644 index 0000000..c61b04a --- /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, mx, my; atol, rtol, check_inferred=false) + + # MinimalSVec + mx = MinimalSVec(x) + my = MinimalSVec(y) + test_rrule(inner, mx, my; atol, rtol, check_inferred=false) +end + +end diff --git a/test/runtests.jl b/test/runtests.jl index 76d08ad..e99badd 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) && isempty(VERSION.prerelease) + println("Testing AD rules") + println("================") + include("chainrules.jl") +end