add chainrules for *-operator

Author: migrosser

Project.toml

@@ -10,6 +10,12 @@ Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
+[weakdeps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+
+[extensions]
+LinearOperatorsChainRulesCoreExt = "ChainRulesCore"
+
 [compat]
 FastClosures = "0.2, 0.3"
 LDLFactorizations = "0.9, 0.10"
@@ -20,6 +26,7 @@ julia = "^1.6.0"
 Arpack = "7d9fca2a-8960-54d3-9f78-7d1dccf2cb97"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 TestSetExtensions = "98d24dd4-01ad-11ea-1b02-c9a08f80db04"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Arpack", "Test", "TestSetExtensions"]
+test = ["Arpack", "Test", "TestSetExtensions", "Zygote"]

ext/LinearOperatorsChainRulesCoreExt.jl

@@ -0,0 +1,51 @@
+module LinearOperatorsChainRulesCoreExt
+
+using LinearOperators
+import ChainRulesCore
+
+function ChainRulesCore.frule((_, Δx, _), ::typeof(*), op::AbstractLinearOperator{T}, x::AbstractVector{S}) where {T, S}
+  y = op*x
+  Δy = op*Δx
+  return y, Δy
+end
+function ChainRulesCore.rrule(::typeof(*), op::AbstractLinearOperator{T}, x::AbstractVector{S}) where {T, S}
+  y = op*x
+  project_x = ChainRulesCore.ProjectTo(x)
+  function mul_pullback(ȳ)
+      x̄ = project_x( adjoint(op)*ChainRulesCore.unthunk(ȳ) )
+      return ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent(), x̄
+  end
+  return y, mul_pullback
+end
+
+function ChainRulesCore.frule((_, Δx, _), ::typeof(*), x::Union{LinearOperators.Adjoint{S, V}, LinearOperators.Transpose{S, V} }, op::AbstractLinearOperator{T}) where {T, S, V <: AbstractVector{S}}
+  y = x*op
+  Δy = Δx*op
+  return y, Δy
+end
+function ChainRulesCore.rrule(::typeof(*), x::LinearOperators.Transpose{S, V}, op::AbstractLinearOperator{T}) where {T, S, V <: AbstractVector{S}}
+  y = x*op
+  project_x = ChainRulesCore.ProjectTo(x)
+  function mul_pullback(ȳ)
+      # needed to make sure that ȳ is recognized as Transposed
+      # ȳ_ = transpose(collect(vec(ChainRulesCore.unthunk(ȳ))))
+      ȳ_ = transpose(vec(ChainRulesCore.unthunk(ȳ)))
+      x̄ = project_x(ȳ_*adjoint(op))
+      return ChainRulesCore.NoTangent(), x̄, ChainRulesCore.NoTangent()
+  end
+  return y, mul_pullback
+end
+function ChainRulesCore.rrule(::typeof(*), x::LinearOperators.Adjoint{S, V}, op::AbstractLinearOperator{T}) where {T, S, V <: AbstractVector{S}}
+  y = x*op
+  project_x = ChainRulesCore.ProjectTo(x)
+  function mul_pullback(ȳ)
+      # needed to make sure that ȳ is recognized as Adjoint
+      # ȳ_ = adjoint(collect(vec(ChainRulesCore.unthunk(ȳ))))
+      ȳ_ = adjoint(conj.(vec(ChainRulesCore.unthunk(ȳ))))
+      x̄ = project_x(ȳ_*adjoint(op))
+      return ChainRulesCore.NoTangent(), x̄, ChainRulesCore.NoTangent()
+  end
+  return y, mul_pullback
+end
+
+end # module

test/runtests.jl

@@ -1,5 +1,6 @@
 using Arpack, Test, TestSetExtensions, LinearOperators
 using LinearAlgebra, SparseArrays
+using Zygote
 include("test_aux.jl")
 
 include("test_linop.jl")
@@ -13,3 +14,4 @@ include("test_callable.jl")
 include("test_deprecated.jl")
 include("test_normest.jl")
 include("test_diag.jl")
+include("test_chainrules.jl")

test/test_chainrules.jl

@@ -0,0 +1,51 @@
+using Zygote
+
+function matmulOp(mat::AbstractArray{T}) where T
+  function prod!(res,x)
+    for i in axes(mat,1)
+      res[i] = transpose(mat[i,:])*x
+    end   
+  end
+
+  function ctprod!(res,x)
+    for i in axes(mat,2)
+      res[i] = dot(mat[:,i],x)
+    end   
+  end
+
+  return LinearOperator{T}(size(mat,1),size(mat,2),false, false, prod!, nothing, ctprod!)
+end
+
+function test_chainrules()
+  @testset ExtendedTestSet "Chainrules" begin
+    for (M,N) in zip([2,3,8,7], [2,4,8,16])
+      for T in [Float64, ComplexF64]
+        mat = simple_matrix(T, M, N)
+        op = matmulOp(mat)
+        x = rand(T,N)
+        xᵀ = transpose(x[1:M])
+        xᴴ = adjoint(x[1:M])
+
+        # test op*x
+        y, g = Zygote.withgradient(v->sum(abs.(op*v)), x)
+        y2, g2 = Zygote.withgradient(v->sum(abs.(mat*v)), x)
+        @test isapprox(y, y2)
+        @test isapprox(g[1], g2[1])
+
+        # test xᵀ*op
+        yt, gt = Zygote.withgradient(v->sum(abs.(v*op)), xᵀ)
+        yt2, gt2 = Zygote.withgradient(v->sum(abs.(v*mat)), xᵀ)
+        @test isapprox(yt, yt2)
+        @test isapprox(gt[1], gt2[1])
+
+        # test xᴴ*op
+        yh, gh = Zygote.withgradient(v->sum(abs.(v*op)), xᴴ)
+        yh2, gh2 = Zygote.withgradient(v->sum(abs.(v*mat)), xᴴ)
+        @test isapprox(yh, yh2)
+        @test isapprox(gh[1], gh2[1])
+      end
+    end
+  end
+end
+
+test_chainrules()