fix: diagm for repeated indices and initial tests

Author: avik-pal

src/stdlibs/LinearAlgebra.jl

@@ -274,9 +274,20 @@ function LinearAlgebra._diagm(
     shape, kv::Pair{<:Integer,<:AnyTracedRArray{T,1}}...
 ) where {T}
     m, n = LinearAlgebra.diagm_size(shape, kv...)
+
+    # For repeated indices we need to aggregate the values
+    kv_updated = Dict{Integer,AnyTracedRArray{T,1}}()
+    for (k, v) in kv
+        if haskey(kv_updated, k)
+            kv_updated[k] = kv_updated[k] + v
+        else
+            kv_updated[k] = v
+        end
+    end
+
     scatter_indices = Matrix{Int64}[]
     concat_inputs = MLIR.IR.Value[]
-    for (k, v) in kv
+    for (k, v) in pairs(kv_updated)
         push!(scatter_indices, diagonal_indices_zero_indexed(m, n, k)[1:length(v), :])
         push!(concat_inputs, get_mlir_data(v))
     end

test/integration/linear_algebra.jl

@@ -130,6 +130,17 @@ end
     @test @jit(diagm(4, 5, x_ra)) ≈ diagm(4, 5, x)
     @test @jit(diagm(6, 6, x_ra)) ≈ diagm(6, 6, x)
     @test_throws DimensionMismatch @jit(diagm(3, 3, x_ra))
+
+    x1 = rand(3)
+    x2 = rand(3)
+    x3 = rand(2)
+    x_ra1 = Reactant.to_rarray(x1)
+    x_ra2 = Reactant.to_rarray(x2)
+    x_ra3 = Reactant.to_rarray(x3)
+
+    @test @jit(diagm(1 => x_ra1)) ≈ diagm(1 => x1)
+    @test @jit(diagm(1 => x_ra1, -1 => x_ra3)) ≈ diagm(1 => x1, -1 => x3)
+    @test @jit(diagm(1 => x_ra1, 1 => x_ra2)) ≈ diagm(1 => x1, 1 => x2)
 end
 
 # TODO: Currently Diagonal(x) * x goes down the generic matmul path but it should clearly be