Make intermediate parameters generated in the for loop accessible for return and generated_quantities

[DominiqueMakowski]

In #2390 (reply in thread), it was suggested the following example to extract the parameters made in the main for loop:

@model function model_Linear(y, x, participant, n)

    # Priors - Fixed Effects
    μ_intercept ~ Normal(0, 0.5)
    μ_x ~ Normal(0, 0.5)

    σ ~ truncated(Normal(0.0, 1), lower=0)

    # Priors - Random Effects
    μ_intercept_random_sd ~ truncated(Normal(0.0, 0.3), lower=0)
    μ_x_random_sd ~ truncated(Normal(0, 0.3), lower=0)

    # Participant-level priors (we only include these if `participant` is not `nothing`)
    if participant !== nothing
        μ_intercept_random ~ filldist(Normal(0, μ_intercept_random_sd), n)
        μ_x_random ~ filldist(Normal(0, μ_x_random_sd), n)
    end

    for i in 1:length(y)
        # Compute μ
        μ = μ_intercept + μ_x * x[i, 1]
        if participant !== nothing
            μ += μ_intercept_random[participant[i]] + μ_x_random[participant[i]] * x[i, 1]
        end

        # Likelihood
        y[i] ~ Normal(μ, σ)
    end

    # For `generated_quantities`
    return (μ_fixed = μ_intercept .+ μ_x .* x[:, 1],)
end

# Convenience method for when participants aren't available.
model_Linear(y, x) = model_Linear(y, x, nothing, nothing)
# Predictions on new data "grid"
grid = DataFrame(x=[0, 50, 100], participant=[missing, missing, missing])
model_predict = model_Linear(fill(missing, 3), grid.x)
pred2 = predict(model_predict, post)
results = generated_quantities(model_predict, post)

In the return (μ_fixed = μ_intercept .+ μ_x .* x[:, 1],), one needs to recompute the value of μ, rather than somehow "extracting" from within the for-loop, which seems to have 2 downsides:

1. Efficiency, if the computation is expensive it seems like an issue to need to redo it to store in for the return
2. Tidyness: in the case of multiple parameters, each having complex computation formulae, it doesn't seem very convenient to duplicate its computation code, once within the for-loop and then for the return.

Is there a workaround that to maybe make available the variables created during the for-loop and return them without explicitly recomputing them?

[torfjelde]

A few things (sorry for being brief; a bit busy atm)

1. Turing.jl will never try to extract standard Julia variables like this automatically. There's just no good way to determine what variables should be included and which should not be included.
2. However, we do offer a syntax for capturing additional variables in the sample call! But this is a very new syntax and so we currently consider it a bit "experimental", though the intention is for this to eventually become a proper feature we also publise:) You can do a := ... where the right-hand side can be whatever value (assuming it can fit in a Chains object). So you could implement the above using μ_tracked[i] = μ on the line below the computation μ = μ_intercept + μ_x * x[i, 1] inside the for-loop (or you just do μ[i] := μ_intercept + μ_x * x[i, 1] and make μ a vector). However, these variables will only be included in sample at the moment, and not in other functions such as predict and generated_quantities 😕
3. Generally speaking, I'm expecting we're not going to do much to tackle the efficiency aspect of it (your (1) point) automatically. More likely, we're going to expose a way for the user to do if !@performing_inference ... end where the ... part can be whatever you want, e.g. an expensive computation that you want to keep track of using := or include in the return. Then it's still up to the user to explicitly handle these things, but this goes with the general philosophy of Turing.jl: don't mess with the user's code except in the fewest of cases when the user explicitly gives us "permission" to do so, i.e. uses ~ or := syntax.
4. It's fairly easy to do this manually atm (and the additional allocations really doens't matter much in this case; negligble compared to the amount of other allocations occuring in the program).

I've made an issue summarising these aspects of := so we can keep track of the progress: #2409

[torfjelde]

To be concrete, using generated_quantities you could do:

@model function model_Linear(y, x, participant, n)

    # Priors - Fixed Effects
    μ_intercept ~ Normal(0, 0.5)
    μ_x ~ Normal(0, 0.5)

    σ ~ truncated(Normal(0.0, 1), lower=0)

    # Priors - Random Effects
    μ_intercept_random_sd ~ truncated(Normal(0.0, 0.3), lower=0)
    μ_x_random_sd ~ truncated(Normal(0, 0.3), lower=0)

    # Participant-level priors (we only include these if `participant` is not `nothing`)
    if participant !== nothing
        μ_intercept_random ~ filldist(Normal(0, μ_intercept_random_sd), n)
        μ_x_random ~ filldist(Normal(0, μ_x_random_sd), n)
    end

    # CHANGED: Pre-allocate the vector.
    μ_fixed = similar(y, eltype(μ_intercept))
    for i in 1:length(y)
        # Compute μ
        μ = μ_intercept + μ_x * x[i, 1]
        # CHANGED: Track.
        μ_fixed[i] = μ
        if participant !== nothing
            μ += μ_intercept_random[participant[i]] + μ_x_random[participant[i]] * x[i, 1]
        end

        # Likelihood
        y[i] ~ Normal(μ, σ)
    end

    # For `generated_quantities`
    # CHANGED: 
    return (μ_fixed = μ_fixed,)
end

Using := you could do

@model function model_Linear(y, x, participant, n)

    # Priors - Fixed Effects
    μ_intercept ~ Normal(0, 0.5)
    μ_x ~ Normal(0, 0.5)

    σ ~ truncated(Normal(0.0, 1), lower=0)

    # Priors - Random Effects
    μ_intercept_random_sd ~ truncated(Normal(0.0, 0.3), lower=0)
    μ_x_random_sd ~ truncated(Normal(0, 0.3), lower=0)

    # Participant-level priors (we only include these if `participant` is not `nothing`)
    if participant !== nothing
        μ_intercept_random ~ filldist(Normal(0, μ_intercept_random_sd), n)
        μ_x_random ~ filldist(Normal(0, μ_x_random_sd), n)
    end

    # CHANGED: Pre-allocate the vector.
    μ_fixed = similar(y, eltype(μ_intercept))
    for i in 1:length(y)
        # Compute μ
        μ = μ_intercept + μ_x * x[i, 1]
        # CHANGED: Use `:=` to track.
        μ_fixed[i] := μ
        if participant !== nothing
            μ += μ_intercept_random[participant[i]] + μ_x_random[participant[i]] * x[i, 1]
        end

        # Likelihood
        y[i] ~ Normal(μ, σ)
    end

    # CHANGED: No need for `return` anymore.
end

And when we finally get something like @is_post_inference (ref: TuringLang/DynamicPPL.jl#714) we could completely avoid the additional allocation during sampling (but, as I said above, this is barely going to affect your inference time so don't think this is something one currently has to worry about):

@model function model_Linear(y, x, participant, n)

    # Priors - Fixed Effects
    μ_intercept ~ Normal(0, 0.5)
    μ_x ~ Normal(0, 0.5)

    σ ~ truncated(Normal(0.0, 1), lower=0)

    # Priors - Random Effects
    μ_intercept_random_sd ~ truncated(Normal(0.0, 0.3), lower=0)
    μ_x_random_sd ~ truncated(Normal(0, 0.3), lower=0)

    # Participant-level priors (we only include these if `participant` is not `nothing`)
    if participant !== nothing
        μ_intercept_random ~ filldist(Normal(0, μ_intercept_random_sd), n)
        μ_x_random ~ filldist(Normal(0, μ_x_random_sd), n)
    end

    # CHANGED: Pre-allocate the vector, only if we're not performing inference.
    if @is_post_inference
        μ_fixed = similar(y, eltype(μ_intercept))
    end
    for i in 1:length(y)
        # Compute μ
        μ = μ_intercept + μ_x * x[i, 1]
        # CHANGED: Use `:=` to track, only if we're not performing inference.
        if @is_post_inference
            μ_fixed[i] := μ
        end
        if participant !== nothing
            μ += μ_intercept_random[participant[i]] + μ_x_random[participant[i]] * x[i, 1]
        end

        # Likelihood
        y[i] ~ Normal(μ, σ)
    end

    # CHANGED: No need for `return` anymore.
end