# Dispatch on the last layer type in Flux

**URL:** <https://discourse.julialang.org/t/dispatch-on-the-last-layer-type-in-flux/82481>\
**Category:** Machine Learning\
**Tags:** question, flux, machine-learning\
**Created:** [June 9, 2022, 9:48am UTC](https://discourse.julialang.org/t/dispatch-on-the-last-layer-type-in-flux/82481 "2022-06-09T09:48:26Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![DoktorMike](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/doktormike/32/2736_2.png) [@DoktorMike](https://discourse.julialang.org/u/DoktorMike)\
**Post date:** [June 9, 2022, 9:48am UTC](https://discourse.julialang.org/t/dispatch-on-the-last-layer-type-in-flux/82481/1 "2022-06-09T09:48:26Z")

</div>

I’m making an implementation of a Evidential Deep Learning model in Flux and I would like to have a function `predict(model, x)` where x is input data. The model is a `Chain` which consists possibly of many layers. I want the predict function to change behavior depending on the type of the **last** layer in the Chain. Is that possible?

Alternatively I could create a struct per model type which contains the Chain and a type and then dispatch on that type. But it seems like there’s probably a better design.

Right now I simply have

```julia
m1 = Chain(Dense(10 => 100, tanh), MyLayerA(100 => 5))
m2 = Chain(Dense(10 => 100, tanh), MyLayerB(100 => 5))

predict(m, x) = m(x)

```

Where you can see that the predict function does not differentiate behavior between the two models. Another way could be to do an if statement in the predict function based on

```julia
typeof(m1.layers[end])

```

but that also seems rather hacky.

---

<div class="post-metadata">

**Author:** ![skleinbo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/skleinbo/32/36080_2.png) [@skleinbo](https://discourse.julialang.org/u/skleinbo)\
**Post date:** [June 9, 2022, 11:00am UTC](https://discourse.julialang.org/t/dispatch-on-the-last-layer-type-in-flux/82481/2 "2022-06-09T11:00:37Z")

</div>

Probably easiest to define `predict(m, x) = predict(typeof(m.layers[end]), m, x)`  
and write dispatches for e.g. `predict(::Type{<:Dense}, m, x) = [...]`

---

<div class="post-metadata">

**Author:** ![DoktorMike](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/doktormike/32/2736_2.png) [@DoktorMike](https://discourse.julialang.org/u/DoktorMike)\
**Post date:** [June 9, 2022, 11:05am UTC](https://discourse.julialang.org/t/dispatch-on-the-last-layer-type-in-flux/82481/3 "2022-06-09T11:05:08Z")

</div>

Cool that looks somewhat like what I was aiming for. 🙏🏻

---

<div class="post-metadata">

**Author:** ![skleinbo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/skleinbo/32/36080_2.png) [@skleinbo](https://discourse.julialang.org/u/skleinbo)\
**Post date:** [June 9, 2022, 11:07am UTC](https://discourse.julialang.org/t/dispatch-on-the-last-layer-type-in-flux/82481/4 "2022-06-09T11:07:01Z")

</div>

I’m almost certain one could make this work without the dynamic dispatch, because `Chain` has all the layer types as parameters in its own type. But I can’t figure out how right now.