# Telling Julia the type of a function's output (and other adventures in type stability)

**URL:** <https://discourse.julialang.org/t/telling-julia-the-type-of-a-functions-output-and-other-adventures-in-type-stability/99681>\
**Category:** New to Julia\
**Created:** [May 31, 2023, 5:08pm UTC](https://discourse.julialang.org/t/telling-julia-the-type-of-a-functions-output-and-other-adventures-in-type-stability/99681 "2023-05-31T17:08:36Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![Bizzi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bizzi/32/51484_2.png) [@Bizzi](https://discourse.julialang.org/u/Bizzi)\
**Post date:** [May 31, 2023, 5:08pm UTC](https://discourse.julialang.org/t/telling-julia-the-type-of-a-functions-output-and-other-adventures-in-type-stability/99681/1 "2023-05-31T17:08:36Z")

</div>

Hey all. I’m having a bit of a hard time understanding types and how to deal with them. Take this small example involving Lux Networks:

```julia
using Lux, Random, JET

input = 1
n = 12
act = tanh
model = 
    Chain(Dense(input => n, act),
          Dense(n => n, act),
          Dense(n => 1), first)
rng = Random.default_rng()
p0, s0 = Lux.setup(rng,model)
x0 = rand(Float32)

u(x,p,st) = x*(model([x],p,st)|>first)

```

Now, if I run `@code_warntype` or `JET.@report_opt` on `model(x0,p0,s0)`, they tell me all is fine. However, I cannot make u, as simple as it is, type stable:

```julia
Arguments
  #self#::Core.Const(u)
  x::Float32
  p::NamedTuple{(:layer_1, :layer_2, :layer_3, :layer_4), Tuple{NamedTuple{(:weight, :bias), Tuple{Matrix{Float32}, Matrix{Float32}}}, NamedTuple{(:weight, :bias), Tuple{Matrix{Float32}, Matrix{Float32}}}, NamedTuple{(:weight, :bias), Tuple{Matrix{Float32}, Matrix{Float32}}}, NamedTuple{(), Tuple{}}}}
  st::Core.Const((layer_1 = NamedTuple(), layer_2 = NamedTuple(), layer_3 = NamedTuple(), layer_4 = NamedTuple()))
Body::Any
1 ─ %1 = Base.vect(x)::Vector{Float32}
│ %2 = Main.model(%1, p, st)::Any
│ %3 = (%2 |> Main.first)::Any
│ %4 = (x * %3)::Any
└── return %4

```

Thus, every other calculation applied to `u` downstream will also be typed as `Any`, which should make the code very slow (right?)

The issue, it seems, is that Julia cannot parse the type of the output of `model` from the types of its inputs, probably because the parameter and state vectors are these large nested structures. That’s ok by me, _but **I** know it_. I could, for example, determine this type by running `model` once on a given set of inputs and then somehow tell the compiler that that’s always going to be the output’s type. However, this doesn’t seem to be possible.

My questions, then, are the following:  
1: How does one deal with the types of function outputs? What should I do when the compiler doesn’t recognize the output’s type for a function I did not write?  
2: Should I really be stressed over this? Does this kind of type stability actually matter for performance? Are there any quick mitigating tricks?

---

<div class="post-metadata">

**Author:** ![GunnarFarneback](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gunnarfarneback/32/1827_2.png) [@GunnarFarneback](https://discourse.julialang.org/u/GunnarFarneback)\
**Post date:** [May 31, 2023, 5:51pm UTC](https://discourse.julialang.org/t/telling-julia-the-type-of-a-functions-output-and-other-adventures-in-type-stability/99681/2 "2023-05-31T17:51:25Z")

</div>

Non-const globals are bad for performance. Try

```julia
const model = 
    Chain(Dense(input => n, act),
          Dense(n => n, act),
          Dense(n => 1), first)

```

---

<div class="post-metadata">

**Author:** ![favba](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/favba/32/2735_2.png) [@favba](https://discourse.julialang.org/u/favba)\
**Post date:** [May 31, 2023, 6:25pm UTC](https://discourse.julialang.org/t/telling-julia-the-type-of-a-functions-output-and-other-adventures-in-type-stability/99681/3 "2023-05-31T18:25:20Z")

</div>

> [@Bizzi](#):
>
> The issue, it seems, is that Julia cannot parse the type of the output of `model` from the types of its inputs, probably because the parameter and state vectors are these large nested structures.

Not really. The problem is that, since `model` is a global variable, at any time you can change it to mean anything else and the function `u` Will have to keep track of it. Therefore “julia” cannot make any prediction of what calling it might return in the future, although it certainly can deduce the type of the `Chain(..` object the variable `model` currently points to.

The comment above fixes the issue. Another approach is to use a let block as

```julia
let const_model=model
   #We define u as global because we want it to be available outside this let scope
    global u(x,p,st) = x(const_model([x],p,st) |> first)
end

```

which makes a local copy of whatever the variable `model` holds **at the moment** of the let block evaluation.

---

<div class="post-metadata">

**Author:** ![Bizzi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bizzi/32/51484_2.png) [@Bizzi](https://discourse.julialang.org/u/Bizzi)\
**Post date:** [May 31, 2023, 6:48pm UTC](https://discourse.julialang.org/t/telling-julia-the-type-of-a-functions-output-and-other-adventures-in-type-stability/99681/4 "2023-05-31T18:48:47Z")

</div>

Oh, I see. So this principle applies even for `model` a function.

Just to be clear, then: The issue here is that the tag `model` was applied to the function returned by `Chain`; When declaring a function the normal way `foo(x)=x` this is not necessary, correct? Even though this definition also lives in the global scope?

---

<div class="post-metadata">

**Author:** ![favba](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/favba/32/2735_2.png) [@favba](https://discourse.julialang.org/u/favba)\
**Post date:** [May 31, 2023, 7:27pm UTC](https://discourse.julialang.org/t/telling-julia-the-type-of-a-functions-output-and-other-adventures-in-type-stability/99681/5 "2023-05-31T19:27:48Z")

</div>

Correct, function names are always constant and cannot be changed:

```julia
julia> my_func(x) = x
my_func (generic function with 1 method)

julia> my_func = 10
ERROR: invalid redefinition of constant my_func
Stacktrace:
 [1] top-level scope
   @ REPL[2]:1

```

I don’t know the internals of `Lux`, but mostly likely the `Chain(..` call returns an object, not an actual function. But it doesn’t matter, as functions are also objects in julia and you can attribute different labels to it (but you cannot change it’s **name** ).

```julia
julia> f = sin
sin (generic function with 14 methods)

julia> f(π/2)
1.0

julia> f = cos
cos (generic function with 14 methods)

julia> f(π/2)
6.123233995736766e-17

julia> sin = cos
ERROR: cannot assign a value to imported variable Base.sin from module Main
Stacktrace:
 [1] top-level scope
   @ REPL[10]:1

```

---

<div class="post-metadata">

**Author:** ![Bizzi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bizzi/32/51484_2.png) [@Bizzi](https://discourse.julialang.org/u/Bizzi)\
**Post date:** [May 31, 2023, 8:23pm UTC](https://discourse.julialang.org/t/telling-julia-the-type-of-a-functions-output-and-other-adventures-in-type-stability/99681/6 "2023-05-31T20:23:07Z")

</div>

Fantastic, thank you very much.