Type instability in `Flux.setup`

```
using Flux

function test_setup(opt, s)
state = Flux.setup(opt, s)
return state
end
s = Chain(
Dense(2 => 100, softsign),
Dense(100 => 2)
)
opt = Adam(0.1)
@code_warntype test_setup(opt, s) # type unstable
```
Output:
```
MethodInstance for GradientFlows.test_setup(::Adam, ::Chain{Tuple{Dense{typeof(softsign), Matrix{Float32}, Vector{Float32}}, Dense{typeof(softsign), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}})
from test_setup(opt, s) @ GradientFlows c:\Users\Math User\.julia\dev\GradientFlows\src\solvers\sbtm.jl:106
Arguments
#self#::Core.Const(GradientFlows.test_setup)
opt::Adam
s::Chain{Tuple{Dense{typeof(softsign), Matrix{Float32}, Vector{Float32}}, Dense{typeof(softsign), Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}}
Locals
state::Any
Body::Any
1 ─ %1 = Flux.setup::Core.Const(Flux.Train.setup)
│ (state = (%1)(opt, s))
└── return state
```
Julia version 1.9.3 and Flux version 0.14.6:
```
(@v1.9) pkg> st Flux
Status `C:\Users\Math User\.julia\environments\v1.9\Project.toml`
[587475ba] Flux v0.14.6
```

Reproduce the provided example with Julia 1.9.3 and Flux 0.14.6, then inspect the Flux.Train.setup entry point involved in the type-unstable call. Use @code_warntype on test_setup to compare the result; done means the returned state is inferred without Any.

Written by the indexing model from the issue text.