Fix multidim solve (#1754)

EnzymeAD · Aug 25, 2024 · 32b7aa2 · 32b7aa2
1 parent 0f7dfcc
commit 32b7aa2
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 3 deletions.
diff --git a/src/internal_rules.jl b/src/internal_rules.jl
@@ -327,6 +327,8 @@ end
     return LinearAlgebra.qr(cache_A, ColumnNorm())
 end
 
+@inline onedimensionalize(::Type{T}) where T <: Array = Vector{eltype(T)}
+
 # y=inv(A) B
 #   dA −= z y^T
 #   dB += z, where  z = inv(A^T) dy
@@ -371,19 +373,19 @@ function EnzymeRules.augmented_primal(config, func::Const{typeof(\)}, ::Type{RT}
 
 @static if VERSION < v"1.8.0"
     UT = Union{
-        LinearAlgebra.Diagonal{eltype(AT), BT},
+        LinearAlgebra.Diagonal{eltype(AT), onedimensionalize(BT)},
         LinearAlgebra.LowerTriangular{eltype(AT), AT},
         LinearAlgebra.UpperTriangular{eltype(AT), AT},
         LinearAlgebra.LU{eltype(AT), AT},
         LinearAlgebra.QRCompactWY{eltype(AT), AT}
     }
 else
     UT = Union{
-        LinearAlgebra.Diagonal{eltype(AT), BT},
+        LinearAlgebra.Diagonal{eltype(AT), onedimensionalize(BT)},
         LinearAlgebra.LowerTriangular{eltype(AT), AT},
         LinearAlgebra.UpperTriangular{eltype(AT), AT},
         LinearAlgebra.LU{eltype(AT), AT, Vector{Int}},
-        LinearAlgebra.QRPivoted{eltype(AT), AT, BT, Vector{Int}}
+        LinearAlgebra.QRPivoted{eltype(AT), AT, onedimensionalize(BT), Vector{Int}}
     }
 end
 

diff --git a/test/internal_rules.jl b/test/internal_rules.jl
@@ -133,6 +133,21 @@ end
 
     y = A \ b
     @test dA ≈ (-z * transpose(y))
+
+    # Ensure multi dim doesn't crash
+    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),
+    )
 end
 
 function tr_solv(A, B, uplo, trans, diag, idx)