Type stability / tuple lengths

Consider the following function that, given a tuple t, computes partial products of the elements specified by a second tuple parts as follows

function partialprods(t::NTuple{N,Int}, parts::NTuple{M,Int}) where {N,M}
    all(0 ≤ i ≤ length(t) for i ∈ parts) || error("Dimension out of range.")
    issorted(parts) || error("`parts` must be sorted in ascending order")
    ntuple(i -> prod(t[parts[i]+1 : parts[i+1]]), Val(M-1))
end

julia> partialprods((1,2,3,4), (0,2,4))
(2, 12)

That is, adjacent elements of parts specify ranges over which to compute the product, and partialprods outputs a tuple containing them.

I observed the following interesting behavior where partialprods is only sometimes type stable (in Julia 1.6.0)

julia> wrapper1() = partialprods((1,2,3,4), (0,2,4))
wrapper1 (generic function with 1 method)

julia> wrapper1()
(2, 12)

julia> @code_warntype wrapper1()
Variables
  #self#::Core.Const(wrapper1)

Body::Tuple{Int64, Int64}
1 ─ %1 = Core.tuple(1, 2, 3, 4)::Core.Const((1, 2, 3, 4))
│   %2 = Core.tuple(0, 2, 4)::Core.Const((0, 2, 4))
│   %3 = Main.partialprods(%1, %2)::Core.Const((2, 12))
└──      return %3

julia> wrapper2() = partialprods((1,2,3,4), (1,3,4))
wrapper2 (generic function with 1 method)

julia> wrapper2()
(6, 4)

julia> @code_warntype wrapper2()
Variables
  #self#::Core.Const(wrapper2)

Body::Tuple{Any, Any}
1 ─ %1 = Core.tuple(1, 2, 3, 4)::Core.Const((1, 2, 3, 4))
│   %2 = Core.tuple(1, 3, 4)::Core.Const((1, 3, 4))
│   %3 = Main.partialprods(%1, %2)::Tuple{Any, Any}
└──      return %3

I would be grateful to anyone who could explain to me what I’ve done that introduces the type instability here. It seems the output type of partialprods should be inferrable from the length of parts, but it is not (always). Note that the first case seems special in that I could not identify another that gave type stable behavior.

t[parts[i]+1 : parts[i+1]] is not type stable, unless the compiler can know ahead of time that this is of constant length as a function of i, like in your type stable example.
Instead of slicing to make tuples of dynamic lengths, use a loop or prod with a function argument to index into t.

Thank you - I see now. Looks like one should generally avoid indexing tuples with non-constant ranges. In my original function, I thought the prod might serve as a barrier since it results, I think, in an Int for any NTuple{N,Int} where N (without a function argument), but it appears that the compiler does not infer this automatically.

I tried your suggestion to use a loop instead, which worked. As an alternative, I replaced

    ntuple(i -> prod(t[parts[i]+1 : parts[i+1]]), Val(M-1))

with

    ntuple(i -> prod(t[j] for j ∈ t[parts[i]+1]:parts[i+1]), Val(M-1))

Thus, “lazy” indexing is the way to go here.