Add ChainRulesCore extension

Project.toml

@@ -1,22 +1,32 @@
 name = "VectorInterface"
 uuid = "409d34a3-91d5-4945-b6ec-7529ddf182d8"
 authors = ["Jutho Haegeman <jutho.haegeman@ugent.be> 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"
 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"]

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

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

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