Fix unflatten_long_ops

SciML · Feb 19, 2021 · 7d5839d · 7d5839d
1 parent ddcba6c
commit 7d5839d
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/src/build_function.jl b/src/build_function.jl
@@ -53,18 +53,18 @@ function build_function(args...;target = JuliaTarget(),kwargs...)
 end
 
 function unflatten_args(f, args, N=4)
-    length(args) < N && return Term{Number}(f, args)
-    unflatten_args(f, [Term{Number}(f, group)
+    length(args) < N && return Term{Real}(f, args)
+    unflatten_args(f, [Term{Real}(f, group)
                                        for group in Iterators.partition(args, N)], N)
 end
 
 function unflatten_long_ops(op, N=4)
     op = value(op)
-    !istree(op) && return
+    !istree(op) && return Num(op)
     rule1 = @rule((+)(~~x) => length(~~x) > N ? unflatten_args(+, ~~x, 4) : nothing)
     rule2 = @rule((*)(~~x) => length(~~x) > N ? unflatten_args(*, ~~x, 4) : nothing)
 
-    Rewriters.Postwalk(Rewriters.Chain([rule1, rule2]))(op)
+    Num(Rewriters.Postwalk(Rewriters.Chain([rule1, rule2]))(op))
 end