Issue using `register_symbolic` with a function that returns `Vector{Float64}` after the 5.3.0 release

[chris-hampel-CA]

Hello,

My project has been stuck on Symbolics 5.2.0 because I ran into an error after trying to upgrade to 5.3.0. It no longer allowed me to use a registered function to return a vector of values to a variable that is defined as a vector.
I realize I am falling behind getting the newest versions of other packages in the SciML ecosystem so I can no longer put off figuring out this issue.

A simple example using Modeling toolkit below functions properly when below using a version of Symbolics less than 5.3.0. The error is due to change in version 5.2.0 to 5.3.0 where Broadcast.materialize was changed to Broadcase.copy
v5.2.0...v5.3.0
due to another issue #882

The dummy script that works for 5.2.0 but fails on 5.3.0

using ModelingToolkit, NonlinearSolve

function vec_return(T, n)
    return [T for _ in 1:n]
end
@register_symbolic vec_return(T, n)::Vector{Float64}

nx=3
Tx=500.0

vars = @variables (x[1:nx], a, b, c)
pars = @parameters (n=nx, T=Tx)

eqs = [
    x .~ vec_return(T, n)
    a ~ 1 + x[1]
    b ~ 2 + x[2]
    c ~ 3 + x[3]
]

@named sys = NonlinearSystem(eqs, [vars...;], pars)
guess = [x[1] => 1, x[2] => 1, x[3] => 1,
        a => 1, b => 1, c => 1]
prob = NonlinearProblem(sys, guess, [])
sol = solve(prob, NewtonRaphson())

The result with Symbolics 5.2.0

julia> sol1 = solve(prob, NewtonRaphson())
u: 6-element Vector{Float64}:
 500.0
 500.0
 500.0
 501.0
 502.0
 503.0

The result with Symbolics 5.3.0

julia> eqs = [
           x .~ vec_return(T, n)
           a ~ 1 + x[1]
           b ~ 2 + x[2]
           c ~ 3 + x[3]
       ]
ERROR: axes of (vec_return(T, n)) not known
Stacktrace:
 [1] error(s::String)
   @ Base .\error.jl:35
 [2] axes(A::Symbolics.Arr{Num, 1})
   @ Symbolics C:\Users\Christopher Hampel\.julia\packages\Symbolics\gy00O\src\arrays.jl:528
 [3] combine_axes
   @ .\broadcast.jl:512 [inlined]
 [4] instantiate
   @ .\broadcast.jl:294 [inlined]
 [5] materialize(bc::Base.Broadcast.Broadcasted{Symbolics.SymWrapBroadcast, Nothing, typeof(~), Tuple{Symbolics.Arr{Num, 1}, Symbolics.Arr{Num, 1}}})
   @ Base.Broadcast .\broadcast.jl:873
 [6] top-level scope
   @ REPL[14]:1

Is there anyway to make a change in my code to allow this or can the previous functionality of 5.2.0 be re-added to Symbolics src code?
Thanks.

[Christopher Rackauckas]

I'm surprised you ever had a vector return work with that. @register_symbolic always assumed a scalar output. You need to update this to @register_array_symbolic for array functions and it should be fine.

[chris-hampel-CA]

Ok I see that is was added in 5.13.0. Thanks for letting me know!

I am trying to test the script on 5.13.0, but I am still getting stumped with figuring out the exact implementation. I think I am closely following the example near line 530 in the MTK test file

If you don't mind taking a peak below and offering some guidance, would be much appreciated.

function vec_return(T, n)
    [T, T, T] 
end
@register_array_symbolic vec_return(T, n) begin
    size=(3,)
    eltype=promote_type(eltype(T), eltype(n)) #=optional?, fails to build eqs if I do not specify this=#
end

nx=3.0 
Tx=500.0

vars = @variables (x[1:3], a, b, c)
pars = @parameters (n=nx, T=Tx)

eqs = [
    x ~ vec_return(T, n)
    a ~ 1 + x[1]
    b ~ 2 + x[2]
    c ~ 3 + x[3]
]
@show eqs

@named sys = NonlinearSystem(eqs, vars, pars)
guess = [x[1] => 1, x[2] => 1, x[3] => 1,
        a => 1, b => 1, c => 1]
prob = NonlinearProblem(sys, guess, [])
sol = solve(prob, NewtonRaphson())

The eqs print as:

eqs = Any[(broadcast(~, x, Main.var"##mock test#292".vec_return(T, n)))[1], (broadcast(~, x, Main.var"##mock test#292".vec_return(T, n)))[2], (broadcast(~, x, Main.var"##mock test#292".vec_return(T, n)))[3], a ~ 1 + x[1], b ~ 2+ x[2], c ~ 3 + x[3]]

The error is:

  LoadError: MethodError: no method matching -(::SymbolicUtils.BasicSymbolic{Any}, ::SymbolicUtils.BasicSymbolic{Real})

Does anything look blatantly incorrect in the implementation when using @register_array_symbolic in your mind? Thanks.