# Broadcasting Flux.jl model over vectors

**URL:** <https://discourse.julialang.org/t/broadcasting-flux-jl-model-over-vectors/83775>\
**Category:** General Usage\
**Tags:** question, broadcast, flux\
**Created:** [July 5, 2022, 1:41pm UTC](https://discourse.julialang.org/t/broadcasting-flux-jl-model-over-vectors/83775 "2022-07-05T13:41:27Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![RomeoV](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/romeov/32/37687_2.png) [@RomeoV](https://discourse.julialang.org/u/RomeoV)\
**Post date:** [July 5, 2022, 1:41pm UTC](https://discourse.julialang.org/t/broadcasting-flux-jl-model-over-vectors/83775/1 "2022-07-05T13:41:27Z")

</div>

Hello,  
Flux.jl based models are usually evaluated in a batch-wise fashion for performance reasons, i.e.

```julia
xs_batch = rand(10, 64); # 64 ten-dimensional features
ys_batch = model(xs_batch)
@assert size(ys_batch) == (1, 64)

```

For cleanly writing a pseudo-code implementation, I would like to define `xs` as an `AbstractVector` of samples which lie continuously in memory, and then write `model.(xs)`. The overloaded implementation should then automatically batch the samples, compute the results batch-by-batch and then unbatch them again to return a `AbstractVector{Output_t}`.

```julia
# I want this
xs = [rand(10) for _ in 1:64]; # or better `rand(10, 64) |> unbatch`
ys = model.(xs) # automatically batches
@assert length(ys) == 64

```

I was able to redefine the broadcasting rule for my model (see below), and `model.(xs)` works.  
But when defining another function `f(x) = 2*model(x)` and then computing `f.(xs)`, the redefined broadcasting rule doesn’t get used anymore…

```julia
f(x) = 2*model(x)
f.(xs) # this doesn't use my custom batching anymore :(

```

Here’s a MWE including the redefined broadcasting for my model.

```julia
using Flux
using MLUtils
using Transducers

Input_t = Vector{Float32}
Output_t = Vector{Float32}
data = rand(Float32, 3, 20) |> unbatch; # makes vector of slices
model = Chain(Dense(3, 100), Dense(100, 100), Dense(100, 2));

Broadcast.broadcasted(model_::typeof(model), xs::AbstractVector{Input_t}) = begin
    batchsize = 16
    model_batched(xs :: AbstractVector{Input_t}) :: AbstractVector{Output_t} = begin
            @info length(xs)
            return model_(xs |> batch) |> unbatch
        end
    return collect(xs |> Partition(batchsize; flush=true) |> MapCat(model_batched)) # or tcollect/dcollect
end

model.(data) # works (!), prints the @info
f(x) = 2*model(x)
f.(data) # also "works", but doesn't use the provided broadcasting implementation

```

Any pointers how to make `f` also use the broadcasting rule would be appreciated!

Cheers

---

<div class="post-metadata">

**Author:** ![albheim](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/albheim/32/34660_2.png) [@albheim](https://discourse.julialang.org/u/albheim)\
**Post date:** [July 5, 2022, 4:16pm UTC](https://discourse.julialang.org/t/broadcasting-flux-jl-model-over-vectors/83775/2 "2022-07-05T16:16:09Z")

</div>

To make `f` use the broadcasting rule you would need it to be the same type, and if you want that to work in general for arbitrary functions then you probably want to dispatch on any callable with your specific data type. This could be done, though I don’t think it is a good idea since it might affect other code in unexpected ways.

It seems like `unbatch` copies the original data, so if you wanted that preserved without copying you might want to create your own implementation for that.  
But then you might also just create a wrapper type for the data and dispatch the broadcast on that type instead of the model type maybe? Seems easier if you dont want to keep the model type constant.

Though I can’t really say I see how this becomes that much cleaner than just being clear about having a matrix where the batching is in one of the dimensions, just seems a bit messier with having to do this broadcast dispatch in some reasonable way 😅

---

<div class="post-metadata">

**Author:** ![bertschi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bertschi/32/33462_2.png) [@bertschi](https://discourse.julialang.org/u/bertschi)\
**Post date:** [July 5, 2022, 9:32pm UTC](https://discourse.julialang.org/t/broadcasting-flux-jl-model-over-vectors/83775/3 "2022-07-05T21:32:41Z")

</div>

If I understand your use case correctly, this would fit very well the idea behind my library [JJ.jl](https://github.com/bertschi/JJ.jl) (inspired by APL/J). Using it, you could define your model/function for a vector input and apply it over a batch by giving it `rank` of one, i.e.,

```julia
using JJ

model = Dense(10, 2)
f(x::AbstractVector) = 2 .* model(x) 

batch = randn(10, 64)
rank"f 1"(batch) |> size # would give (2, 64)

```

This gives indeed very clean code also for more complicated models, e.g., see my [example](https://github.com/bertschi/JJ.jl/blob/main/examples/transformers.jl) of a simple transfomer layer.

Internally, this builds on [JuliennedArrays](https://github.com/bramtayl/JuliennedArrays.jl) which you could also use directly, if a more explicit implementation is desired, i.e.,

```julia
using JuliennedArrays
Align(f.(Slices(batch, 1)), 1)

```

where `Slices` corresponds to your `unbatch` and `Align` reconstructs an array from an array of arrays. In any case, some other libraries such as [SplitApplyCombine](https://github.com/JuliaData/SplitApplyCombine.jl) or the functions `eachslice` and `stack` in Julia 1.9 (nightly) provide similar functionalities.

---

<div class="post-metadata">

**Author:** ![RomeoV](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/romeov/32/37687_2.png) [@RomeoV](https://discourse.julialang.org/u/RomeoV)\
**Post date:** [July 5, 2022, 10:11pm UTC](https://discourse.julialang.org/t/broadcasting-flux-jl-model-over-vectors/83775/4 "2022-07-05T22:11:54Z")

</div>

@albheim  
Regarding the reason why I believe this would make the code more clean, consider testing a property of the model for each sample - e.g. check that the the prediction error is less than `0.5`.  
We can define a function

```julia
function test_abs_error_less_than(x, y, abs_err)
    y_pred = model(x)
    return abs(y - y_pred) <= abs_err
end

```

Then, I would like to execute

```julia
percent_passed = mean(test_abs_error_less_than.(xs, ys, 0.5))

```

This is imo extremely concise and skips any fumbling with defining some dataloader over the xs and so on.  
Note that often we cannot compute `model(xs_full_batch)` because we run out of GPU memory.

Before playing around with this, I thought that broadcasting the `test_abs_error_less_than` function essentially defined a function like this

```julia
function test_abs_error_less_than.(xs, y,s abs_err)
    ys_pred = model.(xs)
    return abs.(ys .- ys_pred) .<= abs_err
end

```

i.e. with all functions transformed to their broadcasted equivalent. But I understand now that this is not the case? This would go against my observations on a simplified test:

```julia
using Test
g(x) = 2*sin(x)
xs = [1., 2., 3.]
@test_throws MethodError g(xs)
g.(xs) # this works

```

I assume that in the call to `g.(xs)`, we implicitly define a function like `g.(xs) = 2.*sin.(x)`. But this seems to be different to my model example above.

---

<div class="post-metadata">

**Author:** ![RomeoV](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/romeov/32/37687_2.png) [@RomeoV](https://discourse.julialang.org/u/RomeoV)\
**Post date:** [July 5, 2022, 10:21pm UTC](https://discourse.julialang.org/t/broadcasting-flux-jl-model-over-vectors/83775/5 "2022-07-05T22:21:50Z")

</div>

@bertschi  
Thanks, this looks very interesting! Perhaps the JuliennedArrays can take care of the extra copies when calling `batch` and `unbatch`. `JJ.jl` also looks very interesting, although I am not very familiar with the J language, so I will have to look into it a bit.

The `eachslice` and `stack` functions are actually a good alternative to the `batch` and `unbatch` functions from MLUtils. In my version of julia (1.6), `batch` is significantly faster than `stack` though. I couldn’t find any details on the changes you mentioned for julia 1.9.

---

<div class="post-metadata">

**Author:** ![bertschi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bertschi/32/33462_2.png) [@bertschi](https://discourse.julialang.org/u/bertschi)\
**Post date:** [July 5, 2022, 11:01pm UTC](https://discourse.julialang.org/t/broadcasting-flux-jl-model-over-vectors/83775/6 "2022-07-05T23:01:11Z")

</div>

The changes for 1.9 are discussed here: [https://github.com/JuliaLang/julia/pull/43334](https://github.com/JuliaLang/julia/pull/43334)

Guess your model of broadcasting might work, but it is not complete, i.e., it explains broadcasting in terms of more broadcasting. I would rather think of it as being a for-loop, i.e.,

```julia
A = randn(2, 3)
g(x) = 2*sin(x)

[g(A[idx]) for idx in CartesianIndices(A)] # roughly same as g.(A)

```

In any case, I had misunderstood your use case. `model.(data)` would handle all inputs sequentially and not take advantage of mini-batching. To that end, I would probably just define a small function instead of trying to overload the broadcasting syntax:

```julia
function minibatched(model, xs::AbstractVector{Input_t}; batchsize = 16)
    xs |> Partition(batchsize; flush = true) |> MapCat(xs -> model(xs |> batch) |> unbatch) |> collect
end

minibatched(model, data) # instead of model.(data)

```

If you urgently need a more concise syntax, you could probably overload one of the currently unused symbols that get parsed as infix operators (see [julia - User-defined infix operator - Stack Overflow](https://stackoverflow.com/questions/60321301/user-defined-infix-operator)):

```julia
⋄(m, xs) = minibatched(m, xs)
model ⋄ data

```