# Help with type instability needed

I would like to write a function `grideval(f,x)` that evaluates a function `f` on the cartesian grid described by the tuple `x` of one-dimensional grid coordinates. This is what I have so far:

``````"""
cartesian(x::AbstractVector...)
cartesian(x::NTuple{N,AbstractVector})

Cartesian product.

# Examples
``
julia> cartesian([1,2],[3,4])
2×2 ApproxTools.TensorGrid{...}:
(1, 3)  (1, 4)
(2, 3)  (2, 4)
``
"""
cartesian(x::AbstractVector...) = cartesian(x)
cartesian(x::NTuple{N,AbstractVector}) where {N} = TensorGrid{Tuple{eltype.(x)...},length(x),typeof(x)}(x)
struct TensorGrid{T,N,F} <: AbstractArray{T,N}
factors::F
end
Base.size(x::TensorGrid) = length.(x.factors)
Base.getindex(x::TensorGrid{<:Any,N,<:Any}, i::Vararg{Int,N}) where {N} = map(getindex,x.factors,i)

"""
grideval(f, x::AbstractVector...)
grideval(f, x::NTuple{N,AbstractVector})

Evaluate `f` on the grid spanned by the `x`.

# Examples
``
julia> grideval(*, ([1,2],[3,4]))
2×2 Array{Int64,2}:
3  4
6  8
``
"""
grideval(f, x::AbstractVector...) = grideval(f,x)
grideval(f, x::NTuple{N,AbstractVector}) where {N} = (x->f(x...)).(cartesian(x))
``````

Unfortunately, `grideval` turns out to have strange type-stability properties. For example, this is fine:

``````julia> function f1()
x = linspace(0,1,10)
grideval((x1,x2)->2*x1*x2, (x,x))
end
@code_warntype f1()
Variables:
#self# <optimized out>
#11::##11#12
x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}

Body:
begin
x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}} = \$(Expr(:invoke, MethodInstance for linspace(::Type{Float64}, ::Int64, ::Int64, ::Int64, ::Int64), :(Base.linspace), :(Base.Float64), 0, 1, 10, 1)) # line 3:
#11::##11#12 = \$(Expr(:new, :(Main.##11#12)))
SSAValue(1) = (Core.tuple)(x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}, x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}})::Tuple{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}
return \$(Expr(:invoke, MethodInstance for grideval(::Function, ::Tuple{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}), :(Main.grideval), :(#11), SSAValue(1)))
end::Array{Float64,2}

``````

But this isn’t:

``````julia> function f2()
a = 2
x = linspace(0,1,10)
grideval((x1,x2)->a*x1*x2, (x,x))
end
@code_warntype f2()
Variables:
#self# <optimized out>
#13::##13#14{Int64}
a <optimized out>
x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}

Body:
begin  # line 3:
x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}} = \$(Expr(:invoke, MethodInstance for linspace(::Type{Float64}, ::Int64, ::Int64, ::Int64, ::Int64), :(Base.linspace), :(Base.Float64), 0, 1, 10, 1)) # line 4:
#13::##13#14{Int64} = \$(Expr(:new, ##13#14{Int64}, 2))
SSAValue(1) = (Core.tuple)(x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}, x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}})::Tuple{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}
return \$(Expr(:invoke, MethodInstance for grideval(::Function, ::Tuple{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}), :(Main.grideval), :(#13), SSAValue(1)))
end::Any
``````

Is there a way to make the second version type-stable as well?

This probably works:

``````function f2()
a = 2
x = linspace(0,1,10)
let a=a
grideval((x1,x2)->a*x1*x2, (x,x))
end
end
``````

1 Like

Unfortunately it doesn’t:

``````julia> @code_warntype f2()
Variables:
#self# <optimized out>
#3::##3#4{Int64}
a@_3 <optimized out>
a@_4 <optimized out>
x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}

Body:
begin  # line 3:
x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}} = \$(Expr(:invoke, MethodInstance for linspace(::Type{Float64}, ::Int64, ::Int64, ::Int64, ::Int64), :(Base.linspace), :(Base.Float64), 0, 1, 10, 1)) # line 5:
#3::##3#4{Int64} = \$(Expr(:new, ##3#4{Int64}, 2))
SSAValue(2) = (Core.tuple)(x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}, x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}})::Tuple{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}
return \$(Expr(:invoke, MethodInstance for grideval(::Function, ::Tuple{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}), :(Main.grideval), :(#3), SSAValue(2)))
end::Any
``````

I am aware of the issues with closures capturing variables, but as far as I can tell they shouldn’t apply here: even without the `let` block, `a` is just a simple local integer variable, and the compiler claims it optimised it away.

It is type stable in Julia 0.7.

For 0.6 you can try this workaround:

``````grideval(f, ::Type{T}, x::NTuple{N,AbstractVector}) where {T,N} = (x->f(x...)::T).(cartesian(x))
function f2()
a = 2
x = linspace(0,1,10)
grideval((x1,x2)->a*x1*x2, Float64, (x,x))
end
``````
2 Likes

For completeness, here’s an alternative approach to `grideval` that is type-stable on 0.6 as well:

``````# Utility function for grideval
# x is only used for its length. It would be cleaner to pass Val(N) instead, but unlike
# Base.tail, Val(N-1) is not type-stable.
shape(x::NTuple{N},i,l) where {N} = (shape(Base.tail(x),i,l)..., i == N ? l : 1)
shape(::NTuple{0},i,l) = ()

``````

For completeness, here’s an alternative approach to `grideval` that is type-stable on 0.6 as well:

``````"""
grideval(f, x::AbstractVector...)
grideval(f, x::NTuple{N,AbstractVector})

Evaluate `f` on the grid spanned by the `x`.

# Examples
``
julia> grideval(*, ([1,2],[3,4]))
2×2 Array{Int64,2}:
3  4
6  8
``
"""
grideval(f, x::AbstractVector...) = grideval(f,x)
@generated grideval(f, x::NTuple{N,AbstractVector}) where {N} = :(Base.Cartesian.@ncall(\$N,broadcast,f,i->reshape(x[i],shape(x,i))))

# Utility function for grideval
# x is only used for its length. It would be cleaner to pass Val(N) instead, but unlike
# Base.tail, Val(N-1) is not type-stable.
shape(x::NTuple{N,Any},i) where {N} = (shape(Base.tail(x),i)..., i == N ? length(x[i]) : 1)
shape(::NTuple{0},i) = ()
``````