时间戳

首页 > 解决方案 > 切片元组时避免类型不稳定

julia - 切片元组时避免类型不稳定

问题描述

这个问题是JuliaLang Zulip 帮助台中出现的问题的改编版本。

假设我有一个函数，该函数接收异构类型Tuple并将返回该元组的一个切片，其中切片索引可以仅从类型信息中静态推断。如何以正确推断输出类型的方式编写函数？

例如，假设我的功能是

function f(t::Tuple, A::Array{T, N}) where {T, N}
    if T <: AbstractFloat
        imin = 1
    elseif T <: Integer
        imin = 2
    else
        imin = 3
    end
    imax = N+2    
    t[imin:imax]
end

我们看到类型推断只计算出这会产生 a Tuple，而不是它的长度或元素类型，即使所有需要的信息在编译时都可用？

julia> let t = (:a, "b", 2, 3.0, Val(1), 2+im), A = rand(Int, 3,3)
           Base.return_types(f, Tuple{typeof(t), typeof(A)})
       end
1-element Array{Any,1}:
 Tuple

我怎样才能写出f这样的作品？

标签： julia

解决方案

我最喜欢的策略（但也许有更简单的方法？）是编写一个@generated函数来手动确保 julia 在编译时执行我想要的类型级别操作：

@generated function f2(t::Tuple, A::Array{T, N}) where {T, N}
    if T <: AbstractFloat
        imin = 1
    elseif T <: Integer
        imin = 2
    else
        imin = 3
    end
    imax = N+2
    out_expr = Expr(:tuple, (:(t[$i]) for i ∈ imin:imax)...)
end

这里的想法是，在生成的函数体中，在编译时，我们确定什么imin和imax是，然后我们手动为函数体构建一个表达式，读取为(t[imin], t[imin+1], ..., t[imax-1], t[imax]).

无论出于何种原因，julia 能够更好地推理序列而getindex(::Tuple, ::Int)不是切片元组，即使是静态已知切片，所以通过手动构建这个表达式，编译器能够做我们想做的事情：

julia> let t = (:a, "b", 2, 3.0, Val(1), 2+im), A = rand(Int, 3,3)
           Base.return_types(f2, Tuple{typeof(t), typeof(A)})
       end
1-element Array{Any,1}:
 Tuple{String,Int64,Float64}

瞧，推断的输出类型是Tuple长度为 3 的 a，其元素静态已知为 a String，Int和Float64！

推荐阅读