Inference regression on master for "collector arguments"

[iamed2]

This shows up as nightly test failures for AxisArrays.jl which were not present on 0.6.

Here is the simplest example:

Indexing: Error During Test
  Test threw an exception of type ErrorException
  Expression: @inferred(D[1, 1, 1, :]) == @inferred(D[1, 1, 1, 1:1]) == @inferred(D[1, 1, 1, [1]]) == AxisArray([10], Axis{:dim_4}(Base.OneTo(1)))
  return type AxisArrays.AxisArray{Int64,1,Array{Int64,1},Tuple{AxisArrays.Axis{:dim_4,Base.OneTo{Int64}}}} does not match inferred return type AxisArrays.AxisArray{Int64,1,Array{Int64,1},_} where _

JuliaArrays/AxisArrays.jl#92

[martinholters]

It's due to a collector argument, but the situation is a bit more complicated because the collector argument is not only used for collecting. It's a bit like the following:

foo(in1::Tuple{}, in2::Tuple{}) = _foo((), in1, in2)
_foo(out, in1::Tuple{}, in2::Tuple{}) = out
_foo(out, in1::Tuple, in2::Tuple{}) = _foo((out..., Val{length(out)}()), Base.tail(in1), in2)
_foo(out, in1::Tuple{}, in2::Tuple) = _foo((out..., 0), in1, Base.tail(in2))
_foo(out, in1::Tuple, in2::Tuple) = _foo((out..., in1[1]), Base.tail(in1), Base.tail(in2))

If it wasn't for the Val{length(out)}() in the second method, one could quite easily rewrite this without a collector argument. Introducing another tuple that only shrinks during the recursion and holds the potentially necessary values does the trick here:

makevals(::Tuple{}) = ()
makevals(in::Tuple) = (makevals(Base.tail(in))..., Val{length(Base.tail(in))}())
foo(in1::Tuple, in2::Tuple) = _foo(makevals(in1), in1, in2)
_foo(vals, in1::Tuple{}, in2::Tuple{}) = ()
_foo(vals, in1::Tuple, in2::Tuple{}) = (vals[1], _foo(Base.tail(vals), Base.tail(in1), in2)...)
_foo(vals, in1::Tuple{}, in2::Tuple) = (0, _foo(vals, in1, Base.tail(in2))...)
_foo(vals, in1::Tuple, in2::Tuple) = (in1[1], _foo(Base.tail(vals), Base.tail(in1), Base.tail(in2))...)

However, if it wasn't a Val here but something which is expensive to construct, this would be rather wasteful. One would want the vals tuple to be replaced with a tuple just holding the types, but just moving the instantiation doesn't work, i.e.

makevals(in::Tuple) = (makevals(Base.tail(in))..., Val{length(Base.tail(in))})
_foo(vals, in1::Tuple, in2::Tuple{}) = (vals[1](), _foo(Base.tail(vals), Base.tail(in1), in2)...)
# other methods as above

again is not inferable, because the return type of makevals is Tuple{DataType, DataType, ...}, not Tuple{Type{Val{0}}, Type{Val{1}}, ...}.

I'm not sure how inferability and avoiding unnecessary instantiations could be achieved at the same time. Any ideas anyone?

[timholy]

Here's an even simpler example, inferrably computing 2^N:

twototheN(ref::Tuple{}) = ()
twototheN(ref::Tuple{Any,Vararg{Any}}) = _twototheN((true, true), Base.tail(ref))
_twototheN(out, ref::Tuple{Any,Vararg{Any}}) = _twototheN((out..., out...), Base.tail(ref))
_twototheN(out, ref::Tuple{}) = out

twototheN((1,2,3)) is inferrable on 0.6 but not on master.

[vtjnash]

arbitrary arithmetic is not expected to be inferrable

[timholy]

This isn't really arithmetic, it's just exercising the type system.

I understand this is a consequence of the redesign of inference to make sure it converges. Wouldn't there be a way to tell that the second argument to twototheN will always get shorter and thus this recursion is finite?

[martinholters]

No, during inference, the second argument could be a Vararg tuple, which wouldn't get shorter.

[timholy]

Since there's a workaround using @pure it's not a big deal.

[martinholters]

If possible, I'd recommend the workaround of not using collector arguments, though:

twototheN(ref::Tuple{}) = (true,)
twototheN(ref::Tuple) = (twototheN(Base.tail(ref))..., twototheN(Base.tail(ref))...)

[vtjnash]

No, during inference, the second argument could be a Vararg tuple, which wouldn't get shorter.

Inference also can't easily determine that this will result in the program terminating. For example, we could just replace the base case with: _foo(out, in1::Tuple{}, in2::Tuple{}) = _foo((), out, out)

We could add heuristics to detect any arbitrary subset of the halting problem, but current experience with base has indicated that our rules are already sufficiently permissive such that there are simple rewrites (for humans, not for computers) that ensure the useful cases are inferrable while the uninferrable / barely inferrable cases can get rejected fairly quickly.

[Jutho]

Is this another instance of the same issue:

tuple_split(t::Tuple, ::Type{Val{N}}) where {N} = _tsplit((), t, Val{N})
@inline _tsplit(t1::NTuple{N1}, t2::NTuple{N2}, ::Type{Val{N1}}) where {N1,N2} = t1, t2
@inline _tsplit(t1::NTuple{N1}, t2::NTuple{N2}, ::Type{Val{N3}}) where {N1,N2,N3}= _tsplit((t1...,t2[1]), Base.tail(t2), Val{N3})

Base.Test.@inferred tuple_split((1,2,3,4,5,6),Val{2}) passes in Julia 0.6 but not in 0.7. I think this qualifies as a useful case? If there is another approach for writing an inferable function for splitting tuples, I'd be interested to know.

[Jutho]

Another one, permuting tuples (where p is assumed to be a valid permutation)

tpermute(t::NTuple{N,T}, p::NTuple{N,Int}) where {T,N} = _tpermute((), t, p)
_tpermute(tdst::NTuple{N,T}, tsrc::NTuple{N,T}, p::NTuple{N,Int}) where {T,N} = tdst
_tpermute(tdst::NTuple{M,T}, tsrc::NTuple{N,T}, p::NTuple{N,Int}) where {T,M,N} = _tpermute((tdst...,tsrc[p[M+1]]), tsrc, p)

On Julia 0.6, the return type of e.g. tpermute((1,2,3,4,5,6),(3,4,6,1,5,2)) is correctly inferred as
NTuple{6,Int64}. On Julia 0.7, the return type seems to be Tuple{Vararg{T,6}} where T<:Int64 lighting up in red in @code_warntype. Shouldn't that just be Tuple{Vararg{Int64,6}} which is the same as NTuple{6,Int64}. Or is that just a display issue?

If type inference does no longer work correctly for these cases, it forces us back to using more @generated functions.

Also strange, on neither Julia 0.6 or 0.7, the call to _tpermute gets inlined. However, adding explicit @inline has effect on Julia0.6 but not on 0.7?

[vtjnash]

Those shouldn't be inferrable, but this is:

tpermute(t::NTuple{N, Any}, p::NTuple{N, Int}) where {T} = ntuple(i -> t[p[i]], Val{N}())

[Jutho]

😳 to much trying lispy style tuple manipulation made me forgot about the easy solution. Thanks for that. Even without ntuple, the simple approach would have been:

tpermute(t::NTuple{N,T}, p::NTuple{N,Int}) where {T,N} = _tpermute(t, p)
_tpermute(t, p::Tuple{}) = ()
_tpermute(t, p::Tuple) = (t[p[1]], _tpermute(t, tail(p))...) 

which also is inferred correctly. But than again, for my overly complicated implementation of above, it's actually not the final tuple length that is incorrectly inferred, but the type of elements, as reported in #23278. I noticed you have already started working on that, so thanks for that.

[martinholters]

Regarding #22513 (comment), this should work:

tuple_split(t::Tuple{Vararg{Any,N}}, ::Type{Val{N}}) where {N} = (t, ())
function tuple_split(t::Tuple, ::Type{Val{N}}) where {N}
    t1, t2 = tuple_split(Base.front(t), Val{N})
    return (t1, (t2..., last(t)))
end