AD support via chainrules (#177)

Co-authored-by: Daniel Karrasch <[email protected]>
Co-authored-by: Jeff Fessler <[email protected]>

Author: 3 people

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"

docs/src/history.md

@@ -1,5 +1,13 @@
 # Version history
 
+## What's new in v3.9
+
+* The application of `LinearMap`s to vectors operation, i.e., `(A,x) -> A*x = A(x)`, is now
+  differentiable w.r.t. to the input `x` for integration with machine learning frameworks
+  such as [`Flux.jl`](https://fluxml.ai/Flux.jl/stable/). The reverse differentiation rule
+  makes `A::LinearMap` usable as a static, i.e., non-trainable, layer in a network, and
+  requires the adjoint `A'` of `A` to be defined.
+
 ## What's new in v3.8
 
 * A new map called [`InverseMap`](@ref) is introduced. Letting an `InverseMap` act on a

src/LinearMaps.jl

@@ -11,6 +11,9 @@ using SparseArrays
 
 import Statistics: mean
 
+using ChainRulesCore: unthunk, NoTangent, @thunk
+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")