# Type instability for simple MLP call using ComponentArrays

**URL:** <https://discourse.julialang.org/t/type-instability-for-simple-mlp-call-using-componentarrays/126191>\
**Category:** General Usage\
**Tags:** zygote, sciml, type-stability, neural-network, autodiff\
**Created:** [February 22, 2025, 4:48pm UTC](https://discourse.julialang.org/t/type-instability-for-simple-mlp-call-using-componentarrays/126191 "2025-02-22T16:48:39Z")\
**Posts on this page:** 3\
**Page:** 1

<div class="post-metadata">

**Author:** ![lasse](https://avatars.discourse-cdn.com/v4/letter/l/96bed5/32.png) [@lasse](https://discourse.julialang.org/u/lasse)\
**Post date:** [February 22, 2025, 4:48pm UTC](https://discourse.julialang.org/t/type-instability-for-simple-mlp-call-using-componentarrays/126191/1 "2025-02-22T16:48:39Z")

</div>

Hi,

Is there a way to make Zygote gradients type stable when optimizing over a ComponentArray of parameters? If not - what is a good alternative that is compatible with the SciML ecosystem?

Using @NamedTuple does as I expect, but is not possible for my problem as it is not a SciMLStructure

```julia
using Lux, Zygote, ComponentArrays, StableRNGs

const model = Chain(
    Dense(2, 8, tanh), 
    Dense(8, 1))
    
ps, st = Lux.setup(StableRNG(1111), model)  

const _st = st
const P = ComponentArray(ps=ps) 

#* This is type stable
@benchmark Zygote.gradient(p -> model(x, p, _st)[1][1], ps) # fast
@code_warntype Zygote.gradient(p -> model(x, p, _st)[1][1], ps)

#* This is NOT type stable for some reason...
@benchmark Zygote.gradient(p -> model(x, p, _st)[1][1], P.ps) # slow
@code_warntype Zygote.gradient(p -> model(x, p, _st)[1][1], P.ps) 

```

Best regards.

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [March 2, 2025, 11:09pm UTC](https://discourse.julialang.org/t/type-instability-for-simple-mlp-call-using-componentarrays/126191/2 "2025-03-02T23:09:52Z")

</div>

The tends to be a byproduct of Zygote. We are moving more and more to Enzyme which tends to be better with inference.

---

<div class="post-metadata">

**Author:** ![lasse](https://avatars.discourse-cdn.com/v4/letter/l/96bed5/32.png) [@lasse](https://discourse.julialang.org/u/lasse)\
**Post date:** [March 6, 2025, 5:17pm UTC](https://discourse.julialang.org/t/type-instability-for-simple-mlp-call-using-componentarrays/126191/3 "2025-03-06T17:17:58Z")

</div>

Ah, okay. That makes sense - thanks!