Error in `gradient` of custom `AbstractMatrix` subtypes

[simonmandlik]

I am struggling with defining custom rrule and computing gradients a custom AbstractMatrix subtype.

For simplicity, assume naive matmul wrapped in a struct

using Flux, Zygote, ChainRulesCore

struct MyMatrix{T, U <: AbstractMatrix{T}} <: AbstractMatrix{T}
    W::U
end

function Base.:*(A::MyMatrix, B::AbstractMatrix)
    C = zeros(size(A.W, 1), size(B, 2))
    for i = 1:size(A.W, 1), j = 1:size(B, 2), k = 1:size(A.W, 2)
        C[i,j] += A.W[i,k] * B[k,j]
    end
    return C
end

function ChainRulesCore.rrule(::typeof(Base.:*), A::MyMatrix, B::AbstractMatrix)
    A.W * B, Δ -> (NoTangent(), NoTangent(), NoTangent())
end

W, B = rand(2,2), rand(2,2)
@assert MyMatrix(W) * B ≈ W * B
gradient(A -> sum(A*B), MyMatrix(W))

The last line gives

julia> gradient(B -> sum(MyMatrix(W)*B), B)
ERROR: type Nothing has no field method
Stacktrace:
  [1] getproperty
    @ ./Base.jl:42 [inlined]
  [2] matching_cr_sig(t::IRTools.Inner.Meta, s::Nothing)
    @ Zygote ~/.julia/packages/Zygote/umM0L/src/compiler/chainrules.jl:72
  [3] has_chain_rrule(T::Type)
    @ Zygote ~/.julia/packages/Zygote/umM0L/src/compiler/chainrules.jl:54
  [4] #s3060#1217
    @ ~/.julia/packages/Zygote/umM0L/src/compiler/interface2.jl:20 [inlined]
  [5] var"#s3060#1217"(::Any, ctx::Any, f::Any, args::Any)
    @ Zygote ./none:0
  [6] (::Core.GeneratedFunctionStub)(::Any, ::Vararg{Any})
    @ Core ./boot.jl:580
  [7] _pullback
    @ ./REPL[8]:1 [inlined]
  [8] _pullback(ctx::Zygote.Context, f::var"#5#6", args::Matrix{Float64})
    @ Zygote ~/.julia/packages/Zygote/umM0L/src/compiler/interface2.jl:0
  [9] _pullback(f::Function, args::Matrix{Float64})
    @ Zygote ~/.julia/packages/Zygote/umM0L/src/compiler/interface.jl:34
 [10] pullback(f::Function, args::Matrix{Float64})
    @ Zygote ~/.julia/packages/Zygote/umM0L/src/compiler/interface.jl:40
 [11] gradient(f::Function, args::Matrix{Float64})
    @ Zygote ~/.julia/packages/Zygote/umM0L/src/compiler/interface.jl:75
 [12] top-level scope
    @ REPL[8]:1
 [13] top-level scope
    @ ~/.julia/packages/CUDA/sCev8/src/initialization.jl:52

What is the correct way to define rrule for Base.:* in situations like this one?

[mcabbott]

If you add some println statements, you can see it never calls your rule. And the reason is that it's ambiguous:

julia> rrule(*, MyMatrix(W), B)
ERROR: MethodError: rrule(::typeof(*), ::MyMatrix{Float64, Matrix{Float64}}, ::Matrix{Float64}) is ambiguous. Candidates:
  rrule(::typeof(*), A::AbstractVecOrMat{<:Union{Real, Complex}}, B::AbstractVecOrMat{<:Union{Real, Complex}}) in ChainRules at /Users/me/.julia/packages/ChainRules/3HAQW/src/rulesets/Base/arraymath.jl:27
  rrule(::typeof(*), A::MyMatrix, B::AbstractMatrix) in Main at REPL[216]:1
Possible fix, define
  rrule(::typeof(*), ::MyMatrix{var"#s3069", U} where {var"#s3069"<:Union{Real, Complex}, U<:AbstractMatrix{var"#s3069"}}, ::AbstractMatrix{T} where T<:Union{Real, Complex})

If you fix that, it works:

julia> function ChainRulesCore.rrule(::typeof(Base.:*), A::MyMatrix, B::AbstractMatrix{<:Real})
           println("my rrule")
           A.W * B, Δ -> (NoTangent(), NoTangent(), NoTangent())
       end

julia> rrule(*, MyMatrix(W), B)
my rrule
([0.736166425259842 0.6843604402987; 0.15875547985974608 0.31206282839297994], var"#262#263"())

julia> Base.size(A::MyMatrix) = size(A.W)

julia> gradient(A -> sum(A*B), MyMatrix(W))
my rrule
(nothing,)

[simonmandlik]

Thank you @mcabbott!

But what if I need a slightly more complex logic in Base.:* and define rrule only for the multiplication part? Say I want to check matrix dimensions before multiplying:

using Flux, Zygote, ChainRulesCore

import Base: *

struct MyMatrix{T, U <: AbstractMatrix{T}} <: AbstractMatrix{T}
    W::U
end

Base.size(A::MyMatrix) = size(A.W)

function _check_mul(A::AbstractMatrix, B::AbstractMatrix)
    if size(A, 2) != size(B, 1)
        throw(DimensionMismatch(""))
    end
end

Base.:*(A::MyMatrix, B::AbstractMatrix{<:Real}) = (_check_mul(A, B); _mul(A, B))

function _mul(A::MyMatrix, B::AbstractMatrix{<:Real})
    C = zeros(size(A.W, 1), size(B, 2))
    for i = 1:size(A.W, 1), j = 1:size(B, 2), k = 1:size(A.W, 2)
        C[i,j] += A.W[i,k] * B[k,j]
    end
    return C
end

function ChainRulesCore.rrule(::typeof(_mul), A::MyMatrix, B::AbstractMatrix{<:Real})
    println("my rrule")
    A.W * B, Δ -> (NoTangent(), NoTangent(), NoTangent())
end

In this case, the custom rrule isn't used:

julia> W, B = rand(2,2), rand(2,2)
([0.944699256606415 0.537992115332553; 0.5155560874253797 0.8018392559893107], [0.7970690302087193 0.5407063289340399; 0.3239745012636194 0.6712300327662797])

julia> gradient(A -> sum(A*B), MyMatrix(W))
ERROR: getindex not defined for MyMatrix{Float64, Matrix{Float64}}
Stacktrace:
  [1] error(::String, ::Type)
    @ Base ./error.jl:42
  [2] error_if_canonical_getindex(::IndexCartesian, ::MyMatrix{Float64, Matrix{Float64}}, ::Int64, ::Int64)
    @ Base ./abstractarray.jl:1231
  [3] matmul2x2!(C::Matrix{Float64}, tA::Char, tB::Char, A::MyMatrix{Float64, Matrix{Float64}}, B::FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool})
    @ LinearAlgebra /Applications/Julia-1.7.app/Contents/Resources/julia/share/julia/stdlib/v1.7/LinearAlgebra/src/matmul.jl:0
  [4] generic_matmatmul!(C::Matrix{Float64}, tA::Char, tB::Char, A::MyMatrix{Float64, Matrix{Float64}}, B::FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool})
    @ LinearAlgebra /Applications/Julia-1.7.app/Contents/Resources/julia/share/julia/stdlib/v1.7/LinearAlgebra/src/matmul.jl:793
...

[mcabbott]

Right, AD never sees _check_mul or _mul, because like mul! they are called after * which it has a rule. They are both run on the forward pass, though. Since your type claims to be a matrix and operates with *, Zygote believes it, and applies what's mathematically true for such objects.

[simonmandlik]

I see, but I can imagine situations where this is really inconvenient. For example, what if I want to implement an AbstractMatrix type where matrix multiplication can be implemented much more efficiently for both forward and backward passes? Something like diagonal matrices, sparse arrays etc.

[mcabbott]

Sometimes this will just work. (::Diagonal) * (::Diagonal) makes another, so the gradient will also be a Diagonal, and so the reverse pass will re-use the same more efficient * method (twice) without you telling AD anything.

If that's not true, you either need to provide a more specific rule, which calls your implementations. Or if your forward pass is itself something AD can handle, then you can use https://juliadiff.org/ChainRulesCore.jl/dev/rule_author/superpowers/opt_out.html .

[simonmandlik]

Thanks for the tips. Is @opt_out supported by Zygote? Adding

ChainRulesCore.@opt_out rrule(::typeof(Base.:*), ::MyMatrix, ::AbstractMatrix{<:Real})

didn't help

[mzgubic]

It should be supported but it hasn't been tested in the wild a lot. Would you be able to put together a minimal example?

[simonmandlik]

Here:

using Flux, Zygote, ChainRulesCore

import Base: *

struct MyMatrix{T, U <: AbstractMatrix{T}} <: AbstractMatrix{T}
    W::U
end

Base.size(A::MyMatrix) = size(A.W)

function _check_mul(A::AbstractMatrix, B::AbstractMatrix)
    if size(A, 2) != size(B, 1)
        throw(DimensionMismatch(""))
    end
end

Base.:*(A::MyMatrix, B::AbstractMatrix{<:Real}) = (_check_mul(A, B); _mul(A, B))
ChainRulesCore.@opt_out rrule(::typeof(Base.:*), ::MyMatrix, ::AbstractMatrix{<:Real})

function _mul(A::MyMatrix, B::AbstractMatrix{<:Real})
    C = zeros(size(A.W, 1), size(B, 2))
    for i = 1:size(A.W, 1), j = 1:size(B, 2), k = 1:size(A.W, 2)
        C[i,j] += A.W[i,k] * B[k,j]
    end
    return C
end

function ChainRulesCore.rrule(::typeof(_mul), A::MyMatrix, B::AbstractMatrix{<:Real})
    println("my rrule")
    A.W * B, Δ -> (NoTangent(), NoTangent(), NoTangent())
end

W, B = rand(2,2), rand(2,2)
@assert MyMatrix(W) * B ≈ W * B
gradient(A -> sum(A*B), MyMatrix(W))

(test_flux) pkg> st
      Status `~/test_flux/Project.toml`
  [d360d2e6] ChainRulesCore v1.12.1
  [587475ba] Flux v0.12.9
  [e88e6eb3] Zygote v0.6.34

[mzgubic]

IIUC the issue you have is that without @opt_out the example works, but is slow. And you want to use @opt_out to just run AD?

This isn't the case with the above snippet, i.e. if I omit the ChainRulesCore.@opt_out rrule(::typeof(Base.:*), ::MyMatrix, ::AbstractMatrix{<:Real}) then it claims the getindex is not defined

julia> gradient(A -> sum(A*B), MyMatrix(W))
ERROR: getindex not defined for MyMatrix{Float64, Matrix{Float64}}
Stacktrace:
  [1] error(::String, ::Type)
    @ Base ./error.jl:42
  [2] error_if_canonical_getindex(::IndexCartesian, ::MyMatrix{Float64, Matrix{Float64}}, ::Int64, ::Int64)
    @ Base ./abstractarray.jl:1231
  [3] matmul2x2!(C::Matrix{Float64}, tA::Char, tB::Char, A::MyMatrix{Float64, Matrix{Float64}}, B::FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool})
    @ LinearAlgebra /Applications/Julia-1.7.app/Contents/Resources/julia/share/julia/stdlib/v1.7/LinearAlgebra/src/matmul.jl:0
  [4] generic_matmatmul!(C::Matrix{Float64}, tA::Char, tB::Char, A::MyMatrix{Float64, Matrix{Float64}}, B::FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}, _add::LinearAlgebra.MulAddMul{true, true, Bool, Bool})
    @ LinearAlgebra /Applications/Julia-1.7.app/Contents/Resources/julia/share/julia/stdlib/v1.7/LinearAlgebra/src/matmul.jl:793
  [5] mul!
    @ /Applications/Julia-1.7.app/Contents/Resources/julia/share/julia/stdlib/v1.7/LinearAlgebra/src/matmul.jl:458 [inlined]
  [6] mul!
    @ /Applications/Julia-1.7.app/Contents/Resources/julia/share/julia/stdlib/v1.7/LinearAlgebra/src/matmul.jl:275 [inlined]
  [7] *
    @ /Applications/Julia-1.7.app/Contents/Resources/julia/share/julia/stdlib/v1.7/LinearAlgebra/src/matmul.jl:153 [inlined]
  [8] #1335
    @ ~/JuliaEnvs/Zygote.jl/dev/ChainRules/src/rulesets/Base/arraymath.jl:37 [inlined]
  [9] unthunk
    @ ~/JuliaEnvs/Zygote.jl/dev/ChainRulesCore/src/tangent_types/thunks.jl:197 [inlined]
 [10] wrap_chainrules_output
    @ ~/JuliaEnvs/Zygote.jl/src/compiler/chainrules.jl:104 [inlined]
 [11] map
    @ ./tuple.jl:223 [inlined]
 [12] wrap_chainrules_output
    @ ~/JuliaEnvs/Zygote.jl/src/compiler/chainrules.jl:105 [inlined]
 [13] ZBack
    @ ~/JuliaEnvs/Zygote.jl/src/compiler/chainrules.jl:204 [inlined]
 [14] Pullback
    @ ./REPL[13]:1 [inlined]
 [15] (::typeof(∂(#7)))(Δ::Float64)
    @ Zygote ~/JuliaEnvs/Zygote.jl/src/compiler/interface2.jl:0
 [16] (::Zygote.var"#57#58"{typeof(∂(#7))})(Δ::Float64)
    @ Zygote ~/JuliaEnvs/Zygote.jl/src/compiler/interface.jl:41
 [17] gradient(f::Function, args::MyMatrix{Float64, Matrix{Float64}})
    @ Zygote ~/JuliaEnvs/Zygote.jl/src/compiler/interface.jl:76
 [18] top-level scope
    @ REPL[13]:1

Is getindex just omitted by mistake from the above or am I misunderstanding the issue?

[simonmandlik]

No, the example doesn't work with or without @opt_out. The error is the same, it wants getindex to be defined, because Zygote calls some generic matrix multiplication instead of rrule defined by me.

In this case, I could of course define getindex easily, but what if such access is expensive and I know of much faster implementation of the backward pass?

[mzgubic]

Oh, I see. Apologies, we should have explained what @opt_out does. It says: ignore the general rrule, and do AD instead through this specific case.

That said, Zygote adjoints will take precedence to rrules, regardless of how loose the signature is.

You will need to define getindex anyway, I don't think it's related to AD in this case:

julia> A = MyMatrix(W)
2×2 MyMatrix{Float64, Matrix{Float64}}:
Error showing value of type MyMatrix{Float64, Matrix{Float64}}:
ERROR: getindex not defined for MyMatrix{Float64, Matrix{Float64}}

[mzgubic]

E.g. defining Base.getindex(mm::MyMatrix, args...; kwargs...) = getindex(mm.W, args...; kwargs...)
will make it work:

julia> @btime gradient(A -> sum(A*B), MyMatrix(W2))
  2.462 μs (35 allocations: 2.23 KiB)
([1.4953375515005742 0.6935637353850378; 1.4953375515005742 0.6935637353850378],)