left division with two matrices

[kpobrien]

Hi, I'm trying to use Enzyme for reverse diff of a function that uses left division. For x = A \ b, I get errors if b is a matrix but no errors if b is a vector. Here is a MWE that illustrates the error.

Using a vector for b runs without error. This is a potential workaround but is not ideal as it would require re-writing other code in a way that doesn't reuse the factorization:

using LinearAlgebra
import Enzyme

function test1!(A)
    A[:,1] .= A \ [1.0,0.0]
    A[:,2] .= A \ [0.0,1.0]
    return nothing
end

A = rand(2,2)
dA = [1.0 0.0; 0.0 0.0]

Enzyme.autodiff(
    Enzyme.Reverse,
    test1!,
    Enzyme.Duplicated(A,dA),
)

and this errors:

using LinearAlgebra
import Enzyme

function test2!(A)
    A .= A \ [1.0 0;0.0 1.0]
    return nothing
end

A = rand(2,2)
dA = [1.0 0.0; 0.0 0.0]

Enzyme.autodiff(
    Enzyme.Reverse,
    test2!,
    Enzyme.Duplicated(A,dA),
)

TypeError: in Diagonal, in V, expected V<:AbstractVector{Float64}, got Type{Matrix{Float64}}

Stacktrace:
  [1] augmented_primal
    @ ~/.julia/packages/Enzyme/XGb4o/src/internal_rules.jl:381 [inlined]
  [2] test2!
    @ ./In[3]:5 [inlined]
  [3] diffejulia_test2__6843wrap
    @ ./In[3]:0
  [4] macro expansion
    @ ~/.julia/packages/Enzyme/XGb4o/src/compiler.jl:7049 [inlined]
  [5] enzyme_call
    @ ~/.julia/packages/Enzyme/XGb4o/src/compiler.jl:6658 [inlined]
  [6] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/XGb4o/src/compiler.jl:6535 [inlined]
  [7] autodiff
    @ ~/.julia/packages/Enzyme/XGb4o/src/Enzyme.jl:320 [inlined]
  [8] autodiff
    @ ~/.julia/packages/Enzyme/XGb4o/src/Enzyme.jl:348 [inlined]
  [9] autodiff(mode::EnzymeCore.ReverseMode{false, EnzymeCore.FFIABI, false, false}, f::typeof(test2!), args::EnzymeCore.Duplicated{Matrix{Float64}})
    @ Enzyme ~/.julia/packages/Enzyme/XGb4o/src/Enzyme.jl:329
 [10] macro expansion
    @ ./timing.jl:279 [inlined]
 [11] top-level scope
    @ ./In[3]:12

Is this a bug or intended behavior? Would you suggest trying to fix this in Enzyme.jl or writing my own rule?

Thanks!

[kpobrien]

The same example with MMatrix from StaticArrays.jl instead of Matrix does not error and may be a workaround:

using LinearAlgebra
import Enzyme, StaticArrays

function test3!(A)
    A .= A \ (StaticArrays.@MMatrix [1.0 0;0.0 1.0])
    return nothing
end

A =  StaticArrays.@MMatrix rand(2,2)
dA = StaticArrays.@MMatrix [1.0 0.0; 0.0 0.0]

Enzyme.autodiff(
    Enzyme.Reverse,
    test3!,
    Enzyme.Duplicated(A,dA),
)

[wsmoses]

looks like something which expected a vector got a matrix, let me investigate

[wsmoses]

Should be solved by https://github.com/EnzymeAD/Enzyme.jl/pull/1754

[kpobrien]

I confirmed your PR #1754 resolves the issue, so I'll close this. Thank you!