🤖 Format .jl files

Author: dpo

ext/LinearOperatorsChainRulesCoreExt.jl

@@ -3,49 +3,71 @@ module LinearOperatorsChainRulesCoreExt
 using LinearOperators
 isdefined(Base, :get_extension) ? (import ChainRulesCore) : (import ..ChainRulesCore)
 
-function ChainRulesCore.frule((_, Δx, _), ::typeof(*), op::AbstractLinearOperator{T}, x::AbstractVector{S}) where {T, S}
-  y = op*x
-  Δy = op*Δx
+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
+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̄
+    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
+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
+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()
+    # 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
+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()
+    # 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
+end # module

src/LinearOperators.jl

@@ -32,7 +32,7 @@ include("deprecated.jl")
 @static if !isdefined(Base, :get_extension)
   import Requires
 end
-   
+
 @static if !isdefined(Base, :get_extension)
   function __init__()
     Requires.@require ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" begin

test/runtests.jl

@@ -14,4 +14,4 @@ include("test_callable.jl")
 include("test_deprecated.jl")
 include("test_normest.jl")
 include("test_diag.jl")
-include("test_chainrules.jl")
+include("test_chainrules.jl")

test/test_chainrules.jl

@@ -1,51 +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   
+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   
+  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!)
+  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 (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 = 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)
+        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ᵀ)
+        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ᴴ)
+        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()
+test_chainrules()