# Is there a way to parallelize portions of code in Flux?

**URL:** <https://discourse.julialang.org/t/is-there-a-way-to-parallelize-portions-of-code-in-flux/101096>\
**Category:** Machine Learning\
**Tags:** flux, machine-learning\
**Created:** [July 2, 2023, 10:13pm UTC](https://discourse.julialang.org/t/is-there-a-way-to-parallelize-portions-of-code-in-flux/101096 "2023-07-02T22:13:09Z")\
**Posts on this page:** 8\
**Page:** 1

<div class="post-metadata">

**Author:** ![josemanuel22](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/josemanuel22/32/20668_2.png) [@josemanuel22](https://discourse.julialang.org/u/josemanuel22)\
**Post date:** [July 2, 2023, 10:13pm UTC](https://discourse.julialang.org/t/is-there-a-way-to-parallelize-portions-of-code-in-flux/101096/1 "2023-07-02T22:13:09Z")

</div>

Hello, I have a code where I want to execute a parallel loop within the code portion responsible for automatic differentiation. This is because the cost function I want to compute relies on certain properties of the estimated and real function that are independent of each other, so I could potentially execute them all in parallel. Is there a way to do this? Right now, when I use @distributed, threads, etc., I get a Zigote error.

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [July 3, 2023, 7:12am UTC](https://discourse.julialang.org/t/is-there-a-way-to-parallelize-portions-of-code-in-flux/101096/2 "2023-07-03T07:12:14Z")

</div>

What kind of Zygote error? Can you post a reproducible example?

---

<div class="post-metadata">

**Author:** ![josemanuel22](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/josemanuel22/32/20668_2.png) [@josemanuel22](https://discourse.julialang.org/u/josemanuel22)\
**Post date:** [July 3, 2023, 10:40am UTC](https://discourse.julialang.org/t/is-there-a-way-to-parallelize-portions-of-code-in-flux/101096/3 "2023-07-03T10:40:37Z")

</div>

Hello, I’m actually starting to feel embarrassed to ask you, I think you’re going to do my homework all by yourself 😉

I’m not sure if you remember the example from last time, but in the end, I did manage to make it differentiate. The code would be something like this.

```plaintext
model = Chain(
    Dense(1 => 10, tanh),
    Dense(10 => 1),
)|> gpu

μ₁=-1; σ₁ = 0.5; 
μ₂=-2; σ₂ = 0.4;
realModel(ϵ) = ϵ < 0.5 ? rand(Normal(μ₁, σ₁)) : rand(Normal(μ₂, σ₂))

#Learning with custom loss
μ = 0; stddev = 1
η = 0.001; num_epochs = 200; n_samples = 10000; K = 10
optim = Flux.setup(Flux.Adam(η), model)
losses = []
@showprogress for epoch in 1:num_epochs
    loss, grads = Flux.withgradient(model) do m
        aₖ = zeros(K+1)
        for _ in 1:n_samples
            x = rand(Normal(μ, stddev), K)
            yₖ = m(x')
            y = realModel(rand(Float64))
            aₖ += generate_aₖ(yₖ, y)
        end
        scalar_diff(aₖ ./ sum(aₖ))
    end
    Flux.update!(optim, model, grads[1])
    push!(losses, loss)
end;

```

In this case, I’m trying to learn a bimodal distribution in one dimension (I only have results showing that it works in one dimension, but it should be true in multiple dimensions as well, although I don’t have the proof yet). We are trying to transform a standard normal distribution (0, 1) into a bimodal distribution (toy example). The inner loop is ‘morally parallelizable’ since the order of simulations and their summations doesn’t matter. Therefore, I would like to execute this in parallel, but when I use @distributed, threads, etc., I get the following Zigote error:

```julia
Can't differentiate foreigncall expression $(Expr(:foreigncall, :(:jl_cpu_wake), Nothing, svec(), 0, :(:ccall))). You might want to check the Zygote limitations documentation. 

```

Let me briefly remind you where each thing comes from.

I have defined a loss function that evaluate the convergence of random generated samples from the model towards a uniform data distribution with respect the real observation.

I generate `K` random hypothetical observations and compare how many of them are smaller than the real data in the training set. In other words, the model generates `K` simulations, and we calculate the number of simulations where the generated data is smaller than the actual data. If the model is well trained, this distribution should converge to a uniform distribution.

`generated_a_k` is construted the following way. However, to simplify, I will just say that executions of `generated_a_k` are independent.

```plaintext
scalar_diff(aₖ) = sum((aₖ .- (1 ./ length(aₖ))) .^2)

"""
    sigmoid(ŷ, y)

    Sigmoid function centered at y.
"""
function sigmoid(ŷ, y)
    return sigmoid_fast.((ŷ-y)*10)
end;

"""
    ψₘ(y, m)

    Bump function centered at m. Implemented as a gaussian function.
"""
function ψₘ(y, m)
    stddev = 0.1
    return exp.((-0.5 .* ((y .- m) ./ stddev) .^ 2))
end

"""
    ϕ(yₖ, yₙ)

    Sum of the sigmoid function centered at yₙ applied to the vector yₖ.
"""
function ϕ(yₖ, yₙ)
    return sum(sigmoid.(yₙ, yₖ))
end;

"""
    γ(yₖ, yₙ, m)
    
    Calculate the contribution of ψₘ ∘ ϕ(yₖ, yₙ) to the m bin of the histogram (Vector{Float}).
"""
function γ(yₖ, yₙ::Float64, m::Int64)
    eₘ(m) = [j == m ? 1.0 : 0.0 for j in 0:length(yₖ)]
    return eₘ(m) * ψₘ(ϕ(yₖ, yₙ), m)
end;

"""
    γ_fast(yₖ, yₙ, m)

Apply the γ function to the given parameters. 
This function is faster than the original γ function because it uses StaticArrays.
However because Zygote does not support StaticArrays, this function can not be used in the training process.
"""
function γ_fast(yₖ, yₙ::Float64, m::Int64)
    eₘ(m) = SVector{length(yₖ)+1, Float64}(j == m ? 1.0 : 0.0 for j in 0:length(yₖ))
    return eₘ(m) * ψₘ(ϕ(yₖ, yₙ), m)
end;

"""
    generate_aₖ(ŷ, y)

    Generate a one step histogram (Vector{Float}) of the given vector ŷ of K simulted observations and the real data y.
    generate_aₖ(ŷ, y) = ∑ₖ γ(ŷ, y, k)
"""
generate_aₖ(ŷ, y::Float64) = sum([γ(ŷ, y, k) for k in 0:length(ŷ)])

```

We need to transform the concept of counting (histogram) into a differentiable operation. To do this, I have done the following. First, I have used a sigmoid operation to check if a fictitious observation (out of the `K` generated) is smaller than the real observation. Obviously, this process can be repeated for the `K` observations by simply summing them (using the `ϕ` function). After this, I want to generate differentiable histogram bins. To achieve this, it is sufficient to use a bump function that sums to nearly 1 when the real observation is greater than the K fictitious observations, and it is zero otherwise. We will have `K+1` bump functions, each centered at `1, 2, …`. With the gamma `γ` function, I simply try to associate the previous result with the vector component `i` , which will represent my histogram bins. Finally, `generated_ak` is nothing more than an interation over all the vector components to generate each of the bins.

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [July 3, 2023, 11:11am UTC](https://discourse.julialang.org/t/is-there-a-way-to-parallelize-portions-of-code-in-flux/101096/4 "2023-07-03T11:11:18Z")

</div>

> [@josemanuel22](#):
>
> Hello, I’m actually starting to feel embarrassed to ask you, I think you’re going to do my homework all by yourself 😉

That’s what I get for answering many AD-related posts I guess 🤣

> [@josemanuel22](#):
>
> when I use @distributed, threads, etc.

Those two are actually [very different](https://docs.julialang.org/en/v1/manual/parallel-computing/), not sure if you’re clear on the specific behavior of each one. I think in your case multithreading (with shared memory) is more appropriate, unless the inner iterations are really long and you’re working on a cluster of several machines.

> [@josemanuel22](#):
>
> The inner loop is ‘morally parallelizable’ since the order of simulations and their summations doesn’t matter.

True, but if you naively add a `@threads` to your code as written you’re gonna make it incorrect due to race conditions in `aₖ += ...`.  
In any case, the most straightforward fix is probably using [ThreadsX.jl](https://github.com/tkf/ThreadsX.jl) to parallelize the sum:

```julia
aₖ = ThreadsX.sum(1:n_samples) do _
    x = rand(Normal(μ, stddev), K)
    yₖ = m(x')
    y = realModel(rand(Float64))
    generate_aₖ(yₖ, y)
end

```

Not tested but that’s the gist

> [@josemanuel22](#):
>
> I get the following Zigote error:

I would like to see the actual code you ran in parallel cause I’m a bit surprised by the “foreign call” error.

---

<div class="post-metadata">

**Author:** ![josemanuel22](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/josemanuel22/32/20668_2.png) [@josemanuel22](https://discourse.julialang.org/u/josemanuel22)\
**Post date:** [July 3, 2023, 3:08pm UTC](https://discourse.julialang.org/t/is-there-a-way-to-parallelize-portions-of-code-in-flux/101096/5 "2023-07-03T15:08:51Z")

</div>

> [@gdalle](#):
>
> Those two are actually [very different](https://docs.julialang.org/en/v1/manual/parallel-computing/), not sure if you’re clear on the specific behavior of each one. I think in your case multithreading (with shared memory) is more appropriate, unless the inner iterations are really long and you’re working on a cluster of several machines.

I understand that the difference is the same as between processes and threads in Linux. You’re right that launching threads should be lighter and have the advantage of shared resources.

> [@gdalle](#):
>
> True, but if you naively add a `@threads` to your code as written you’re gonna make it incorrect due to race conditions in `aₖ += ...`.

You’re right, I was thinking that since I didn’t care about the order of the addition, there shouldn’t be any race conditions. But you’re right that the addition should be an atomic operation to avoid any problems.

> [@gdalle](#):
>
> ```julia
> aₖ = ThreadsX.sum(n_samples) do _
> x = rand(Normal(μ, stddev), K)
> yₖ = m(x')
> y = realModel(rand(Float64))
> generate_aₖ(yₖ, y)
> end
> 
> ```

I’ve tried what you mentioned. I assume you meant to say that:

```Julia
aₖ = ThreadsX.sum(n_samples) do _
    x = rand(Normal(μ, stddev), K)
    yₖ = m(x')
    y = realModel(rand(Float64))
    generate_aₖ(yₖ, y)
end

```

Interestingly, I’m getting the following error.

```julia
nested task error: `llvmcall` must be compiled to be called
Stacktrace:

```

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [July 3, 2023, 3:14pm UTC](https://discourse.julialang.org/t/is-there-a-way-to-parallelize-portions-of-code-in-flux/101096/6 "2023-07-03T15:14:47Z")

</div>

> [@josemanuel22](#):
>
> I’ve tried what you mentioned. I assume you meant to say that:

You copy-pasted the same code, but yeah inside the sum I meant `1:n_samples`

---

<div class="post-metadata">

**Author:** ![josemanuel22](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/josemanuel22/32/20668_2.png) [@josemanuel22](https://discourse.julialang.org/u/josemanuel22)\
**Post date:** [July 3, 2023, 3:22pm UTC](https://discourse.julialang.org/t/is-there-a-way-to-parallelize-portions-of-code-in-flux/101096/7 "2023-07-03T15:22:00Z")

</div>

Sorry, you are right

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [July 6, 2023, 4:17pm UTC](https://discourse.julialang.org/t/is-there-a-way-to-parallelize-portions-of-code-in-flux/101096/8 "2023-07-06T16:17:35Z")

</div>

In theory, Zygote (the AD backend Flux uses by default) supports the lower level built-in threading constructs. See [Zygote.jl/test/threads.jl at 2f4937096ee1db4b5a67c1c31fe3ebeab1c96c8c · FluxML/Zygote.jl · GitHub](https://github.com/FluxML/Zygote.jl/blob/2f4937096ee1db4b5a67c1c31fe3ebeab1c96c8c/test/threads.jl) for a couple of examples. I do not believe it supports `@threads` because that adds a bunch of extra function calls, but it looks like you can get away with not using that for your model.