AD support via chainrules

Author: mloubout

Project.toml

@@ -3,9 +3,11 @@ uuid = "7a12625a-238d-50fd-b39a-03d52299707e"
 version = "3.6.1"
 
 [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

@@ -10,6 +10,8 @@ using SparseArrays
 
 import Statistics: mean
 
+import ChainRulesCore: rrule, frule, unthunk, @thunk, NoTangent
+
 using Base: require_one_based_indexing
 
 abstract type LinearMap{T} end
@@ -268,6 +270,7 @@ include("kronecker.jl") # Kronecker product of linear maps
 include("fillmap.jl") # linear maps representing constantly filled matrices
 include("conversion.jl") # conversion of linear maps to matrices
 include("show.jl") # show methods for LinearMap objects
+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::AbstractVecOrMat)
+    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::AbstractVecOrMat)
+    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,7 @@
 [deps]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Quaternions = "94ee1d12-ae83-5a48-8b1c-48b8ff168ae0"

test/rrules.jl

@@ -0,0 +1,21 @@
+using Flux
+
+@testset "AD rules" begin
+    A = LinearMap(rand(10, 10))
+    x = randn(10)
+    g1 = A'*A*x
+    # Multiplication rule
+    g2 = gradient(x -> .5*norm(A*x)^2, x)
+    @test g1 ≈ g2[1]
+    # Call rule
+    g3 = gradient(x -> .5*norm(A(x))^2, x)
+    @test g1 ≈ g3[1]
+
+    g1 = A*A'*x
+    # Multiplication rule
+    g2 = gradient(x -> .5*norm(A'*x)^2, x)
+    @test g1 ≈ g2[1]
+    # Call rule
+    g3 = gradient(x -> .5*norm(A'(x))^2, x)
+    @test g1 ≈ g3[1]
+end

test/runtests.jl

@@ -33,3 +33,5 @@ include("left.jl")
 include("fillmap.jl")
 
 include("nontradaxes.jl")
+
+include("rrule.jl")