# How to write type stable code with branches, where one branch calls a complicated function?

**URL:** <https://discourse.julialang.org/t/how-to-write-type-stable-code-with-branches-where-one-branch-calls-a-complicated-function/102958>\
**Category:** General Usage\
**Tags:** type-stability\
**Created:** [August 18, 2023, 4:09pm UTC](https://discourse.julialang.org/t/how-to-write-type-stable-code-with-branches-where-one-branch-calls-a-complicated-function/102958 "2023-08-18T16:09:58Z")\
**Posts on this page:** 15\
**Page:** 1

<div class="post-metadata">

**Author:** ![tictaccat](https://avatars.discourse-cdn.com/v4/letter/t/9f8e36/32.png) [@tictaccat](https://discourse.julialang.org/u/tictaccat)\
**Post date:** [August 18, 2023, 4:09pm UTC](https://discourse.julialang.org/t/how-to-write-type-stable-code-with-branches-where-one-branch-calls-a-complicated-function/102958/1 "2023-08-18T16:09:58Z")

</div>

I’m having a hard time writing a type-stable function where one branch calls a complicated function and another one doesn’t. Here’s a minified example:

```julia
using Distributions

function f(μ::Real, σ::Real, a::Real)
    if isnan(μ) || isnan(σ)
        return 0.0
    else
        return pdf(Normal(μ, σ), a) 
    end
end

```

Don’t ask why I want to return 0 when mu or sigma are NaN! This is a minified example, but hopefully the semantics I’m trying to achieve in this case are clear even if unmotivated.

I would like `f` to be type stable. However, you’ve probably spotted an issue: the 0.0 is a `Float64`. So a call such as `f(3.0f32, 4.0f32, 1.0f32)` would not be type stable. I could do a heuristic like `zero(promote_type(typeof(μ), typeof(σ)))` in the first branch. But this means I basically have to reverse engineer the promotion behaviour of `pdf` and `Normal`! (To see why it’s hard: do you know if the heuristic I just wrote is correct? Maybe the type of `a` factors in too?) It’s hard to scale this to more complicated functions.

I would just like to tell the compiler: do the reasonable thing here… whatever type the other branch would normally be, make the first branch a zero of that type, so that `f` can be type stable when possible. Is there a way to express this ergonomically?

---

<div class="post-metadata">

**Author:** ![mikmoore](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mikmoore/32/31109_2.png) [@mikmoore](https://discourse.julialang.org/u/mikmoore)\
**Post date:** [August 18, 2023, 4:51pm UTC](https://discourse.julialang.org/t/how-to-write-type-stable-code-with-branches-where-one-branch-calls-a-complicated-function/102958/2 "2023-08-18T16:51:08Z")

</div>

This is a weak point in Julia’s goal of type-generic code.

If the cost of `pdf` is small or the NaN branch will be taken rarely, you could do something simple like

```julia
function f(μ::Real, σ::Real, a::Real)
    p = pdf(Normal(μ, σ), a)
    if isnan(μ) || isnan(σ)
        return zero(p)
    else
        return p
    end
end

```

I didn’t bother to load up `Distributions.jl` for this, but when I used `muladd(μ, σ, a)` instead of `pdf` and looked at the `@code_native`, it had actually optimized itself so that it only computed `p` when it really needed the value. For the zero branch, it simply returned zero of the appropriate type _without_ computing the result of the function call. The same _might_ happen in your example, depending on how much the compiler can infer about `pdf`.

Otherwise this gets more difficult, relying on introspection. `Base.return_types` _may_ be helpful, but really it’s not a function that should be relied upon (it’s also not public API so may break in the future) and it requires some gymnastics to use in a type-stable way. Hopefully someone else can offer something better than that.

---

<div class="post-metadata">

**Author:** ![tictaccat](https://avatars.discourse-cdn.com/v4/letter/t/9f8e36/32.png) [@tictaccat](https://discourse.julialang.org/u/tictaccat)\
**Post date:** [August 18, 2023, 6:57pm UTC](https://discourse.julialang.org/t/how-to-write-type-stable-code-with-branches-where-one-branch-calls-a-complicated-function/102958/3 "2023-08-18T18:57:42Z")

</div>

Indeed it’s tricky! I guess it’s been discussed plenty how easy it is to write non-type-stable code, I hope that some of the discussion about better static systems in Julia can lead to a solution to this problem 🙂

---

<div class="post-metadata">

**Author:** ![zgornel](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/zgornel/32/217487_2.png) [@zgornel](https://discourse.julialang.org/u/zgornel)\
**Post date:** [August 18, 2023, 7:09pm UTC](https://discourse.julialang.org/t/how-to-write-type-stable-code-with-branches-where-one-branch-calls-a-complicated-function/102958/4 "2023-08-18T19:09:01Z")

</div>

Perhaps using `ifelse` may help

---

<div class="post-metadata">

**Author:** ![jbytecode](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jbytecode/32/17719_2.png) [@jbytecode](https://discourse.julialang.org/u/jbytecode)\
**Post date:** [August 18, 2023, 7:33pm UTC](https://discourse.julialang.org/t/how-to-write-type-stable-code-with-branches-where-one-branch-calls-a-complicated-function/102958/5 "2023-08-18T19:33:23Z")

</div>

Type-stability is not possible in a strong definition (at least as it is in Rust) because Julia is a dynamic language. If I am not getting wrong, how about adding some constraints:

```julia
using Distributions

function f(μ::T, σ::T, a::T)::T where {T <: Real}
    if isnan(μ) || isnan(σ)
        return zero(T)
    else
        return pdf(Normal(μ, σ), a) 
    end
end

```

supposing all of the parameters in same type.

---

<div class="post-metadata">

**Author:** ![gdalle](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/gdalle/32/27854_2.png) [@gdalle](https://discourse.julialang.org/u/gdalle)\
**Post date:** [August 18, 2023, 7:34pm UTC](https://discourse.julialang.org/t/how-to-write-type-stable-code-with-branches-where-one-branch-calls-a-complicated-function/102958/6 "2023-08-18T19:34:44Z")

</div>

I will just add that type-genericity in Distributions.jl is still very much a work in progress: see [Fitting does not respect type parameters · Issue #1544 · JuliaStats/Distributions.jl · GitHub](https://github.com/JuliaStats/Distributions.jl/issues/1544) for an example

---

<div class="post-metadata">

**Author:** ![mkitti](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/mkitti/32/12459_2.png) [@mkitti](https://discourse.julialang.org/u/mkitti)\
**Post date:** [August 19, 2023, 1:51am UTC](https://discourse.julialang.org/t/how-to-write-type-stable-code-with-branches-where-one-branch-calls-a-complicated-function/102958/7 "2023-08-19T01:51:52Z")

</div>

You could just `convert` and type assert.

```julia
using Distributions

function f(μ::Real, σ::Real, a::Real)::Float64
    if isnan(μ) || isnan(σ)
        return 0.0
    else
        return convert(Float64, pdf(Normal(μ, σ), a))::Float64
    end
end

```

---

<div class="post-metadata">

**Author:** ![DNF](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dnf/32/10191_2.png) [@DNF](https://discourse.julialang.org/u/DNF)\
**Post date:** [August 19, 2023, 7:36am UTC](https://discourse.julialang.org/t/how-to-write-type-stable-code-with-branches-where-one-branch-calls-a-complicated-function/102958/8 "2023-08-19T07:36:28Z")

</div>

But that just turns everything into Float64. I doubt that is desirable, in the general case. What if the input is a dual number or a complex number or a `BigFloat`, for example? (The example specifies `Real` inputs, but the question is, I believe, intended to broader than the specific example.)

---

<div class="post-metadata">

**Author:** ![Vasily\_Pisarev](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/vasily_pisarev/32/7929_2.png) [@Vasily\_Pisarev](https://discourse.julialang.org/u/Vasily_Pisarev)\
**Post date:** [August 19, 2023, 9:34am UTC](https://discourse.julialang.org/t/how-to-write-type-stable-code-with-branches-where-one-branch-calls-a-complicated-function/102958/9 "2023-08-19T09:34:16Z")

</div>

Seconding @jbytecode.  
A pattern I often have is

```julia
function f(μ::T, σ::T, a::T) where {T<:Real}
    ...
end

function f(μ::Real, σ::Real, a::Real)
    f(promote(μ, σ, a)...)
end

```

---

<div class="post-metadata">

**Author:** ![jishnub](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jishnub/32/33620_2.png) [@jishnub](https://discourse.julialang.org/u/jishnub)\
**Post date:** [September 20, 2023, 12:16pm UTC](https://discourse.julialang.org/t/how-to-write-type-stable-code-with-branches-where-one-branch-calls-a-complicated-function/102958/10 "2023-09-20T12:16:55Z")

</div>

The promote pattern only really works with numbers, and usually won’t work with other types such as `AbstractVector`s

---

<div class="post-metadata">

**Author:** ![abraemer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/abraemer/32/51403_2.png) [@abraemer](https://discourse.julialang.org/u/abraemer)\
**Post date:** [September 20, 2023, 12:25pm UTC](https://discourse.julialang.org/t/how-to-write-type-stable-code-with-branches-where-one-branch-calls-a-complicated-function/102958/11 "2023-09-20T12:25:05Z")

</div>

I’d like to add that “small type instabilities” (that is `Union{..}` of only a few types) are not that bad for performance since Julia has union-splitting (see this [excellent Blog post by @tim.holy](https://julialang.org/blog/2018/08/union-splitting/) for explanation). So it is not necessary to try to make functions 100% type stable all of the time.

---

<div class="post-metadata">

**Author:** ![jishnub](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jishnub/32/33620_2.png) [@jishnub](https://discourse.julialang.org/u/jishnub)\
**Post date:** [September 20, 2023, 1:36pm UTC](https://discourse.julialang.org/t/how-to-write-type-stable-code-with-branches-where-one-branch-calls-a-complicated-function/102958/12 "2023-09-20T13:36:10Z")

</div>

Not having concrete type-inference may complicate downstream inference, though.

```julia
julia> f(x) = x > 1 ? 0.0 : x # a simple function that returns one of two types
f (generic function with 1 method)

julia> @code_typed f(1)
CodeInfo(
1 ─ %1 = Base.slt_int(1, x)::Bool
└── goto #3 if not %1
2 ─ return 0.0
3 ─ return x
) => Union{Float64, Int64}

julia> @code_typed (() -> [f(i) for i in 1:1])()
CodeInfo(
1 ─ goto #3 if not true
2 ─ nothing::Nothing
3 ┄ %3 = $(Expr(:foreigncall, :(:jl_alloc_array_1d), Vector{Int64}, svec(Any, Int64), 0, :(:ccall), Vector{Int64}, 1, 1))::Vector{Int64}
│ Base.arrayset(true, %3, 1, 1)::Vector{Int64}
│ %5 = invoke Base.collect_to!(%3::Vector{Int64}, $(QuoteNode(Base.Generator{UnitRange{Int64}, var"#8#10"}(var"#8#10"(), 1:1)))::Base.Generator{UnitRange{Int64}, var"#8#10"}, 2::Int64, 1::Int64)::Union{Vector{Int64}, Vector{Real}}
└── goto #4
4 ─ return %5
) => Union{Vector{Int64}, Vector{Real}}

```

Perhaps the `Vector{Real}` may be narrowed to `Vector{Float64}` with some careful promotion, but this may fail to return concrete element types in general. There are also cases where one might desire concrete type-inference to avoid any allocation.

---

<div class="post-metadata">

**Author:** ![abraemer](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/abraemer/32/51403_2.png) [@abraemer](https://discourse.julialang.org/u/abraemer)\
**Post date:** [September 20, 2023, 2:04pm UTC](https://discourse.julialang.org/t/how-to-write-type-stable-code-with-branches-where-one-branch-calls-a-complicated-function/102958/13 "2023-09-20T14:04:30Z")

</div>

I agree that this is bad if you put values into containers which then need to have abstract types.  
I am not sure what you mean with :

> [@jishnub](#):
>
> There are also cases where one might desire concrete type-inference to avoid any allocation.

Doing

```julia
julia> f(x) = x > 5 ? 0.0 : x # type unstable
julia> g(::Float64) = "float"
julia> g(::Int) = "int"
julia> h(x) = g(f(x))
julia> using BenchmarkTools; @btime h(x) setup=(x=rand(1:10);)
  2.305 ns (0 allocations: 0 bytes)

```

So no allocations. Union-splitting is not to be confused with dynamic dispatch where Julia needs to do dispatch entirely at runtime. For union-splitting all possible types and the corresponding methods are fully inferred and simply the correct one is selected by a conditional. In the simple example above, each of the calls to `g` is even fully inlined:

```julia
julia> @code_native debuginfo=:none dump_module=:false h(5)
	.text
	pushq	%rbp
	movq	%rsp, %rbp
	cmpq	$6, %rdi
	movabsq	$140073470456400, %rcx # imm = 0x7F656572F250
	movabsq	$140073446523592, %rax # imm = 0x7F656405C2C8
	cmovlq	%rcx, %rax
	popq	%rbp
	retq

```

---

<div class="post-metadata">

**Author:** ![jishnub](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jishnub/32/33620_2.png) [@jishnub](https://discourse.julialang.org/u/jishnub)\
**Post date:** [September 20, 2023, 2:34pm UTC](https://discourse.julialang.org/t/how-to-write-type-stable-code-with-branches-where-one-branch-calls-a-complicated-function/102958/14 "2023-09-20T14:34:15Z")

</div>

Sorry, I think I was still thinking of containers, as is the case when multiple values are returned from a function. E.g.:

```julia
julia> f(i) = i > 0 ? i : 0.0
f (generic function with 1 method)

julia> g(i) = f(i), nothing
g (generic function with 1 method)

julia> @btime g(i) setup=(i=rand(1:4))
  74.393 ns (1 allocation: 16 bytes)
(2, nothing)

```

This is because

```julia
julia> @code_typed g(1)
CodeInfo(
1 ─ %1 = Base.slt_int(0, i)::Bool
└── goto #3 if not %1
2 ─ goto #4
3 ─ goto #4
4 ┄ %5 = φ (#2 => i, #3 => 0.0)::Union{Float64, Int64}
│ %6 = Core.tuple(%5, Main.nothing)::Tuple{Union{Float64, Int64}, Nothing}
└── return %6
) => Tuple{Union{Float64, Int64}, Nothing}

```

where the return type is a `Tuple` containing a `Union`, and not a `Union` of simple `Tuple`s.

---

<div class="post-metadata">

**Author:** ![tom-plaa](https://avatars.discourse-cdn.com/v4/letter/t/d26b3c/32.png) [@tom-plaa](https://discourse.julialang.org/u/tom-plaa)\
**Post date:** [September 20, 2023, 4:04pm UTC](https://discourse.julialang.org/t/how-to-write-type-stable-code-with-branches-where-one-branch-calls-a-complicated-function/102958/15 "2023-09-20T16:04:05Z")

</div>

I have multiple situations like this in code that I write for stochastic simulation and inference. There’s a hack that I use because it makes it so everything is type stable for ForwardDiff. I select one of the arguments of the function and use its type as a parameter. I make sure it’s a argument that will be autodiff’d. I guess in your case I would hack it around like this:

```julia
using Distributions

function f(μ::T, σ::Real, a::Real) where T<:Real
    res = zero(T)
    if !(isnan(μ) || isnan(σ))
        res += pdf(Normal(μ, σ), a) 
    end
    return res
end

```

@Vasily_Pisarev 's way of doing this is the correct one (imo) but I am often too lazy to apply it everywhere. I should probably refactor some more stuff…