why is my withgradient type unstable ?

[filchristou]

I noticed the withgradient function introduces some type instability in a hot loop of mine.
I got it to the following MWE.

using Flux

function unstablewithgradients(model, obs, labels)
    loss, grads = Flux.withgradient(
        model
    ) do m
            internfunc(m, obs, labels)
        end
    return loss, grads
end

# obviously type stable function
function internfunc(m, x, y)
    modelvals = m.(x)
    meanvals = reduce(vcat, modelvals)
    Flux.mse(meanvals, y)
end

then run

fc = Chain(Dense(5=>3, relu), Dense(3=>1))
obs = [fill(5f0, 5) for _ in 1:10]
labels = fill(2f0, 10)

@code_warntype unstablewithgradients(fc, obs, labels)

I am on Flux v0.14.15, and on Julia 1.10.4

[filchristou]

Closing because answered in https://discourse.julialang.org/t/flux-chain-type-unstable-when-broadcasting-inside-gradient/115741

Do you mind if I add a sentence in http://fluxml.ai/Flux.jl/stable/performance/#Evaluate-batches-as-Matrices-of-features about avoiding vector of vectors ?