fix: input permutation to lu (#1302)

* fix: input permutation to lu

* test: add 3d test

Author: avik-pal

src/stdlibs/LinearAlgebra.jl

@@ -594,7 +594,7 @@ function _lu_overload(
 ) where {T,N}
     # TODO: don't ignore the check and allowsingular flags
     # Batching here is in the last dimensions. `Ops.lu` expects the last dimensions
-    permdims = vcat(Int64[N - 1, N], collect(Int64, 1:(N - 2)))
+    permdims = vcat(collect(Int64, 3:N), 1, 2)
     A = Ops.transpose(materialize_traced_array(A), permdims)
     factors, ipiv, perm, info = Reactant.Ops.lu(A)
 

test/integration/linear_algebra.jl

@@ -370,4 +370,13 @@ end
             @test @jit(solve_with_lu(A_ra, B_ra)) ≈ solve_with_lu_batched(A, B)
         end
     end
+
+    @testset "Input Permutation" begin
+        A = rand(Float32, 10, 10, 32)
+        B = rand(Float32, 10, 32)
+        A_ra = Reactant.to_rarray(A)
+        B_ra = Reactant.to_rarray(B)
+
+        @test @jit(solve_with_lu(A_ra, B_ra)) ≈ solve_with_lu_batched(A, B)
+    end
 end