Fix stack overflow on sum of large number of EinExprs

bsc-quantic · Feb 9, 2024 · b5df61a · b5df61a
1 parent ef1192f
commit b5df61a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/SizedEinExpr.jl b/src/SizedEinExpr.jl
@@ -57,9 +57,9 @@ parsuminds(sexpr::SizedEinExpr) = parsuminds(sexpr.path)
 Base.sum!(sexpr::SizedEinExpr, inds) = sum!(sexpr.path, inds)
 Base.sum(sexpr::SizedEinExpr, inds) = sum(sexpr.path, inds)
 
-function Base.sum(sexpr::Vector{SizedEinExpr{L}}; skip = L[]) where {L}
+function Base.sum(sexpr::Vector{SizedEinExpr{L}}; skip=L[]) where {L}
     path = sum(map(x -> x.path, sexpr); skip)
-    size = @compat(allequal(Iterators.map(x -> x.size, sexpr))) ? first(sexpr).size : merge(map(x -> x.size, sexpr)...)
+    size = @compat(allequal(Iterators.map(x -> x.size, sexpr))) ? first(sexpr).size : mapreduce(x -> x.size, merge!, sexpr; init=Dict{L,Int}())
     # size = merge(map(x -> x.size, sexpr)...)
     SizedEinExpr(path, size)
 end