feat: add similarto keyword to build_function

JuliaSymbolics · Jan 29, 2025 · 061625a · 061625a
1 parent 4d6af94
commit 061625a
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/src/build_function.jl b/src/build_function.jl
@@ -288,6 +288,7 @@ function _build_function(target::JuliaTarget, rhss::AbstractArray, args...;
                        outputidxs=nothing,
                        skipzeros = false,
                        force_SA = false,
+                       similarto = nothing,
                        wrap_code = (nothing, nothing),
                        fillzeros = skipzeros && !(rhss isa SparseMatrixCSC),
                        states = LazyState(),
@@ -301,7 +302,9 @@ function _build_function(target::JuliaTarget, rhss::AbstractArray, args...;
     dargs = map((x) -> destructure_arg(x[2], !checkbounds,
                                   Symbol("ˍ₋arg$(x[1])")), enumerate([args...]))
     i = findfirst(x->x isa DestructuredArgs, dargs)
-    similarto = force_SA ? SArray : i === nothing ? Array : dargs[i].name
+    if similarto === nothing
+        similarto = force_SA ? SArray : i === nothing ? Array : dargs[i].name
+    end
 
     oop, iip = iip_config
     oop_body = if oop

diff --git a/test/build_function.jl b/test/build_function.jl
@@ -281,3 +281,12 @@ end
     T = value(x .^ 2)
     @test_nowarn toexpr(T, NameState())
 end
+
+@testset "`similarto` keyword argument" begin
+    @variables x[1:2]
+    T = collect(value(x .^ 2))
+    fn = build_function(T, collect(x); expression = false)[1]
+    @test_throws MethodError fn((1.0, 2.0))
+    fn = build_function(T, collect(x); similarto = Array, expression = false)[1]
+    @test fn((1.0, 2.0)) ≈ [1.0, 4.0]
+end