Julia Broadcasting Issue with SubArray in Complex Expressions

General Usage
, ,
1

I’m encountering an issue with Julia broadcasting when working with SubArrays in complex nested expressions. When using @views to create SubArrays and then applying broadcasting with @. to complex expressions, I get type mismatch errors.

Minimal Reproducible Example

using SymbolicUtils.Code
using SymbolicUtils
using Symbolics
using RuntimeGeneratedFunctions
RuntimeGeneratedFunctions.init(@__MODULE__)

@variables a b c d
@variables p1 p2 p3

function build_f1(input_names, param_names, outputs, exprs)
    def_calls1 = [:(@views $nm = i[$idx]) for (idx, nm) in enumerate(input_names)]
    def_calls2 = [:(@views $nm = p[$idx]) for (idx, nm) in enumerate(param_names)]
    compute_calls = [:($nm = @. $expr) for (nm, expr) in zip(outputs, exprs)]
    return_calls = :(return [$(outputs...)])
    return :(function (i, p)
        $(def_calls1...)
        $(def_calls2...)
        $(compute_calls...)
        $(return_calls)
    end)
end

input_names = [:a, :b]
param_names = [:p1, :p2, :p3]
outputs = [:c, :d]
exprs = toexpr.([a * p1 + b * p2, sin(c * p1) + tanh(b * p3)])
f1_expr = build_f1(input_names, param_names, outputs, exprs)
bf1 = @RuntimeGeneratedFunction(f1_expr)

Print the Generated Function

RuntimeGeneratedFunction(#=in Main=#, #=using Main=#, :((i, p)->begin
          #= e:\JlCode\HydroModels\dev\bug_submit.jl:18 =#
          #= e:\JlCode\HydroModels\dev\bug_submit.jl:19 =#
          #= e:\JlCode\HydroModels\dev\bug_submit.jl:14 =# @views a = i[1]
          #= e:\JlCode\HydroModels\dev\bug_submit.jl:14 =# @views b = i[2]
          #= e:\JlCode\HydroModels\dev\bug_submit.jl:20 =#
          #= e:\JlCode\HydroModels\dev\bug_submit.jl:15 =# @views p1 = p[1]
          #= e:\JlCode\HydroModels\dev\bug_submit.jl:15 =# @views p2 = p[2]
          #= e:\JlCode\HydroModels\dev\bug_submit.jl:15 =# @views p3 = p[3]
          #= e:\JlCode\HydroModels\dev\bug_submit.jl:21 =#
          c = #= e:\JlCode\HydroModels\dev\bug_submit.jl:16 =# @__dot__((+)((*)(a, p1), (*)(b, p2)))
          d = #= e:\JlCode\HydroModels\dev\bug_submit.jl:16 =# @__dot__((+)((tanh)((*)(b, p3)), (sin)((*)(c, p1))))
          #= e:\JlCode\HydroModels\dev\bug_submit.jl:22 =#
          return [c, d]
      end))

Error Message

When I try to use this function with SubArrays, I get an error like:

ERROR: MethodError: no method matching tanh(::Vector{Float64})
The function `tanh` exists, but no method is defined for this combination of argument types.

Observations

  1. The code works fine with simple expressions but fails with complex nested expressions.
  2. Direct use of broadcasting on the expressions works
@views input1 = eachslice(input, dims=1)[1]
@views input2 = eachslice(input, dims=1)[2]
eval(:(@__dot__((+)((tanh)((*)(input1, 2)), (sin)((*)(input2, 3))))))

Question

How can I properly apply broadcasting to complex expressions involving SubArrays in generated code? Is there a way to ensure that all operations in nested expressions are properly broadcasted?

2

I’m confused by the generated function; it seems to me that the views there don’t do anything, because the indexing is scalar. If i and p are arrays, you need : (colon) to get all elements along a certain dimension. Scalar indexing of a multidimensional array indexes the flattened array. I’m also pretty sure that eachslice creates views by default.

What’s also curious is that the error mentions a Vector, while you say the input consists of SubArrays.

Can you provide the example input with which you run your code?

3

  1. First, I used @views to enable the generated function to support CUDA computation, because if a scalar is extracted from a CuArray, it would result in an error;

  2. Then, I used broadcast to support computations between vectors and scalars. Before this, I directly used the generate function feature in the SymbolicUtils.Code module, but it seems to only support computations between scalars. Of course, it’s possible to perform broadcast computation on the generated function fun.(input, params), but this involves repeatedly extracting data a = i[1]…, which makes the computation and gradient calculation very inefficient. So, I wanted to customize a way to generate functions, mainly to support computations between vectors and scalars.

  3. I’m very sorry for forgetting to upload the test data. In fact, the test data is just simple matrices or vectors, see below:

# vector and scalar
input = ones(2, 10)
params = [2, 3, 4]

# scalar and scalar
input = [2, 2]
params = [2, 3, 4]

Finally, thank you for your reply, and after supplementing the data, I hope to learn about a solution from you.