Type inference through an if statement on a type parameter

I am trying to make a function that specializes on the type parameter and Julia seems unable to concretely infer the types. Here is a MWE:

function h(a::NTuple{dim,Int}) where {dim}
       b = dim == 2 ? 2. : "a"
       c = b*b
       return c
end
@code_warntype h((1,1)) # Infers `b` to be a Union{Float64, String}

I ended up using a generated function to work around this as follows:

@generated function m(a::NTuple{dim,Int}) where {dim}
      b = dim == 2 ? 2. : "a"
      c = b*b
      return :($c)
end
@code_warntype h((1,1,1)) # Works fine

Is this an expected behaviour?

Use dim === 2 (THis will be fixed by #24362)

3 Likes

Works on nightly, thanks!