AD support via chainrules

Project.toml

@@ -1,11 +1,13 @@
 name = "LinearMaps"
 uuid = "7a12625a-238d-50fd-b39a-03d52299707e"
-version = "3.8.0"
+version = "3.9.0"
 
 [deps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [compat]
+ChainRulesCore = "1"
 julia = "1.6"

src/LinearMaps.jl

@@ -11,6 +11,9 @@ using SparseArrays
 
 import Statistics: mean
 
+using ChainRulesCore
+import ChainRulesCore: rrule
+
 using Base: require_one_based_indexing
 
 abstract type LinearMap{T} end
@@ -347,6 +350,7 @@ include("conversion.jl") # conversion of linear maps to matrices
 include("show.jl") # show methods for LinearMap objects
 include("getindex.jl") # getindex functionality
 include("inversemap.jl")
+include("chainrules.jl") # AD rules through ChainRulesCore
 
 """
     LinearMap(A::LinearMap; kwargs...)::WrappedMap

src/chainrules.jl

@@ -0,0 +1,19 @@
+function rrule(::typeof(*), A::LinearMap, x::AbstractVector)
+    y = A*x
+    function pullback(dy)
+        DY = unthunk(dy)
+        # Because A is an abstract map, the product is only differentiable w.r.t the input
+        return NoTangent(), NoTangent(), @thunk(A' * DY)
+    end
+    return y, pullback
+end
+
+function rrule(A::LinearMap, x::AbstractVector)
+    y = A*x
+    function pullback(dy)
+        DY = unthunk(dy)
+        # Because A is an abstract map, the product is only differentiable w.r.t the input
+        return NoTangent(), @thunk(A' * DY)
+    end
+    return y, pullback
+end

test/Project.toml

@@ -1,6 +1,8 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -12,5 +14,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 [compat]
 Aqua = "0.5"
 BlockArrays = "0.16"
+ChainRulesCore = "1"
+ChainRulesTestUtils = "1.9"
 Quaternions = "0.5"
 julia = "1.6"

test/rrules.jl

@@ -0,0 +1,16 @@
+using Test, LinearMaps, ChainRulesTestUtils
+using ChainRulesCore: NoTangent
+
+@testset "AD rules" begin
+    x = randn(10)
+    for A in (
+        LinearMap(rand(10, 10)),
+        LinearMap(cumsum, reverse∘cumsum∘reverse, 10),
+        LinearMap((y, x) -> cumsum!(y, x), (y, x) -> reverse!(cumsum!(y, reverse!(copyto!(y, x)))), 10)
+    )
+        test_rrule(*, A ⊢ NoTangent(), x)
+        test_rrule(A ⊢ NoTangent(), x)
+        test_rrule(*, A' ⊢ NoTangent(), x)
+        test_rrule(A' ⊢ NoTangent(), x)
+    end
+end

test/runtests.jl

@@ -39,3 +39,5 @@ include("embeddedmap.jl")
 include("getindex.jl")
 
 include("inversemap.jl")
+
+include("rrules.jl")