Idiomatic evaluate for scalar, allocate and broadcast for array?

General Usage
, ,
1

In several cases I’ve come across the following pattern:

function f(x; kwargs...)
#intense pre-computation that depend on kwargs but not x
result = 0.0 * x; # one way to preallocate for e.g. x::Vector, but not for scalars
for i in contributions 
    if isa(x, Number)
         result += @. (complicated expression involving x, i, and the precomputations)
    else
        @. result += (same expression)
    end
end
return result
end

The main issue is that @. result += ... errors if result is a scalar. The intense precomputations make externally mapping f costly. One solution is to dispatch f(x::Number; kwargs...) = f([x]; kwargs...), but that array allocation looks awkard to me. The one decent solution I can think of is putting the loop over contributions behind a function barrier, e.g. return sumofcontributions(x,...), but that in my cases requires passing many variables from f’s precomputations to sumofcontributions, which also feels inelegant.

What would be the nominal idiomatic pattern here?

2

Have you tried to convert x into a 1-element tuple (x,) instead of a vector? Broadcasting works with tuples, too, and you can also mix tuples and vectors. Alternatively, you could try a fixed-length vector, for example SVector from StaticArrays.

EDIT: I see that you need result to be mutable. So maybe x can be made into an SVector and result into a (mutable) MVector. If you then return SVector(result), chances are good that there will be no allocation.

3

How are you broadcasting non-scalars to a scalar to add to the result? If they’re all scalars, you wouldn’t need the broadcasting to begin with.

Why not index the element out of the mutable result to return the scalar?

1 Like
4

@matthias314 I’d expect that creating an MVector of length 1 would still allocate within f? My other hesitation is that spelling out a type for result (MVector or otherwise) in the body of f also seems not particularly idiomatic.

5

I believe this is something BangBang.jl is designed to handle. I don’t know much about it, but maybe give it a look.

6

How about something like this? (it may be what you had described with f(x::Number; kwargs...) = f([x]; kwargs...) and I misunderstood :slight_smile: )

julia> precomp(;kwargs...)=(sleep(1);return :s)
precomp (generic function with 1 method)

julia> g(x,pre)=(x,pre)

julia> f(x;kwargs...)=(pre=precomp(;kwargs...);g.(x,Ref(pre)))
f (generic function with 1 method)

julia> @time f(1,x=0)
  1.001090 seconds (10 allocations: 416 bytes)
(1, :s)

julia> @time f([1,2],x=0)
  1.001146 seconds (7 allocations: 288 bytes)
2-element Vector{Tuple{Int64, Symbol}}:
 (1, :s)
 (2, :s)

p.s: I am no authority on what is idiomatic…

7

Not necessarily if the MVector is not returned as such:

julia> using StaticArrays

julia> f(x) = (v = MVector(x); v .*= 2; v[1])
f (generic function with 1 method)

julia> @allocated f(1.0)
0

Something else: Can’t you define the complicated expression as a function g(x) within f (so that you have access to f’s local variables) and then use g either with or without broadcasting?

1 Like
8

If any of these captured variables were reassigned, they’ll get boxed and allocated in the current implementation even if they aren’t reassigned after the closure is instantiated.

MVector{1, typeof(x)}(x) is type-generic and statically knowable from specialization over the argument x. Note that StaticArrays.jl really only supports isbits elements; for example, you could instantiate MVector{1, BigInt}, but setindex! would error.