Improve type inference in iterating over `Eigen`

General Usage Performance
1

In the following example:

julia> @code_warntype (A -> ((λ,v) = eigen(A); (λ,v)))(rand(2,2))
MethodInstance for (::var"#19#20")(::Matrix{Float64})
  from (::var"#19#20")(A) @ Main REPL[14]:1
Arguments
  #self#::Core.Const(var"#19#20"())
  A::Matrix{Float64}
Locals
  @_3::Val{:vectors}
  v::Union{Matrix{ComplexF64}, Matrix{Float64}}
  λ::Union{Vector{ComplexF64}, Vector{Float64}}
Body::Tuple{Union{Vector{ComplexF64}, Vector{Float64}}, Union{Matrix{ComplexF64}, Matrix{Float64}}}
1 ─ %1 = Main.eigen(A)::Union{Eigen{ComplexF64, ComplexF64, Matrix{ComplexF64}, Vector{ComplexF64}}, Eigen{Float64, Float64, Matrix{Float64}, Vector{Float64}}}
│   %2 = Base.indexed_iterate(%1, 1)::Union{Tuple{Vector{ComplexF64}, Val{:vectors}}, Tuple{Vector{Float64}, Val{:vectors}}}
│        (λ = Core.getfield(%2, 1))
│        (@_3 = Core.getfield(%2, 2))
│   %5 = Base.indexed_iterate(%1, 2, @_3)::Union{Tuple{Matrix{ComplexF64}, Val{:done}}, Tuple{Matrix{Float64}, Val{:done}}}
│        (v = Core.getfield(%5, 1))
│   %7 = Core.tuple(λ, v)::Tuple{Union{Vector{ComplexF64}, Vector{Float64}}, Union{Matrix{ComplexF64}, Matrix{Float64}}}
└──      return %7

The type of the result is inferred to be a Tuple of Unions, however, I know that this should ideally be a Union of Tuples. Any suggestions on how to nudge the compiler to infer this as a simpler union, aside from adding type-assertions?

2

Union of which Tuples, exactly? Because it’s very possible for a Tuple of Unions to be equivalent to a Union of Tuples e.g. Tuple{Union{Int,String}} == Union{Tuple{Int}, Tuple{String}}

3

In this problem, I know that λ and v have the same eltype, but this information is lost if I iterate over the struct. One may see this by defining

julia> valsvecs(E::LinearAlgebra.Eigen) = (E.values, E.vectors)
valsvecs (generic function with 1 method)

julia> @code_warntype (A -> valsvecs(eigen(A)))(rand(2,2))
MethodInstance for (::var"#5#6")(::Matrix{Float64})
  from (::var"#5#6")(A) @ Main REPL[8]:1
Arguments
  #self#::Core.Const(var"#5#6"())
  A::Matrix{Float64}
Body::Union{Tuple{Vector{ComplexF64}, Matrix{ComplexF64}}, Tuple{Vector{Float64}, Matrix{Float64}}}
1 ─ %1 = Main.eigen(A)::Union{Eigen{ComplexF64, ComplexF64, Matrix{ComplexF64}, Vector{ComplexF64}}, Eigen{Float64, Float64, Matrix{Float64}, Vector{Float64}}}
│   %2 = Main.valsvecs(%1)::Union{Tuple{Vector{ComplexF64}, Matrix{ComplexF64}}, Tuple{Vector{Float64}, Matrix{Float64}}}
└──      return %2

In this, the result is inferred as a Union of Tuples of arrays that have the same eltype. As soon as I iterate over this, though, the link between the eltypes of λ and v is lost. I was wondering if there’s a way to preserve this.

1 Like
4 
julia> function valsvecs2(E::LinearAlgebra.Eigen)
         (λ,v) = E
         (λ,v)
       end

julia> @code_warntype (A -> valsvecs2(eigen(A)))(rand(2,2))
MethodInstance for (::var"#7#8")(::Matrix{Float64})
  from (::var"#7#8")(A) in Main at REPL[9]:1
Arguments
  #self#::Core.Const(var"#7#8"())
  A::Matrix{Float64}
Body::Union{Tuple{Vector{ComplexF64}, Matrix{ComplexF64}}, Tuple{Vector{Float64}, Matrix{Float64}}}
1 ─ %1 = Main.eigen(A)::Union{Eigen{ComplexF64, ComplexF64, Matrix{ComplexF64}, Vector{ComplexF64}}, Eigen{Float64, Float64, Matrix{Float64}, Vector{Float64}}}
│   %2 = Main.valsvecs2(%1)::Union{Tuple{Vector{ComplexF64}, Matrix{ComplexF64}}, Tuple{Vector{Float64}, Matrix{Float64}}}
└──      return %2

Looks like a case of function barriers. eigen(A) is inferred as a Union of a complex Eigen or a float Eigen. If you extract the values and vectors in the same function call, that function will be dispatched on the Eigen and match the types of the values and vectors. If you extract the values and vectors in separate function calls, the caller cannot assume they return the same type. The second example can have the same problem if the function barrier is omitted:

julia> @code_warntype (A -> (E = eigen(A); (E.values, E.vectors)))(rand(2,2))
MethodInstance for (::var"#11#12")(::Matrix{Float64})
  from (::var"#11#12")(A) in Main at REPL[11]:1
Arguments
  #self#::Core.Const(var"#11#12"())
  A::Matrix{Float64}
Locals
  E::Union{Eigen{ComplexF64, ComplexF64, Matrix{ComplexF64}, Vector{ComplexF64}}, Eigen{Float64, Float64, Matrix{Float64}, Vector{Float64}}}
Body::Tuple{Union{Vector{ComplexF64}, Vector{Float64}}, Union{Matrix{ComplexF64}, Matrix{Float64}}}
1 ─      (E = Main.eigen(A))
│   %2 = Base.getproperty(E, :values)::Union{Vector{ComplexF64}, Vector{Float64}}
│   %3 = Base.getproperty(E, :vectors)::Union{Matrix{ComplexF64}, Matrix{Float64}}
│   %4 = Core.tuple(%2, %3)::Tuple{Union{Vector{ComplexF64}, Vector{Float64}}, Union{Matrix{ComplexF64}, Matrix{Float64}}}
└──      return %4

I would’ve expected this to be handled by Union-splitting, I wonder if the compiler could be improved to do this…if there’s ever a github issue, this could serve as a good example.

5

Is there any way to add a type assertion to restore this link between eltypes? Something like this doesn’t seem to work:

(A -> (t = valsvecs2(eigen(A)); (t[1], t[2]::Matrix{eltype(t[1])})))(rand(2,2))

since the eltype of t[1] is not inferred concretely.

6

valsvecs or valsvecs2 return a tuple matching the types, you wouldn’t need to make another tuple. Type assertions (a type annotating instances) occur at runtime, so t[1] is concrete at that point. It’s variable type declarations (a type annotating the left hand of an assignment) which provide extra information at compile-time.