What is the preferred way to iterate through the Cartesian power of an iterable only available at runtime?

Hi,
Sometimes in my code I need to iterate through the Cartesian power of a vector. After searching on the web I can come up with two approaches.

  1. Using product and repeated from Base.Iterators. The problem is that this approach produces type-unstable code.
    Example:
using Base.Iterators

function test_base_iterators(n, k)
    x = collect(1:n)
    for p in product(repeated(x, k)...)
        println(p)
    end
end

Output of @code_warntype:

julia> @code_warntype test_base_iterators(3, 2)
MethodInstance for test_base_iterators(::Int64, ::Int64)
  from test_base_iterators(n, k) in Main at REPL[1]:1
Arguments
  #self#::Core.Const(test_base_iterators)
  n::Int64
  k::Int64
Locals
  @_4::Union{Nothing, Tuple{Any, Union{Bool, Tuple}}}
  x::Vector{Int64}
  p::Any
Body::Nothing
1 ─ %1  = (1:n)::Core.PartialStruct(UnitRange{Int64}, Any[Core.Const(1), Int64])
│         (x = Main.collect(%1))
│   %3  = Main.repeated(x, k)::Base.Iterators.Take{Base.Iterators.Repeated{Vector{Int64}}}
│   %4  = Core._apply_iterate(Base.iterate, Main.product, %3)::Base.Iterators.ProductIterator
│         (@_4 = Base.iterate(%4))
│   %6  = (@_4 === nothing)::Bool
│   %7  = Base.not_int(%6)::Bool
└──       goto #4 if not %7
2 ┄ %9  = @_4::Tuple{Any, Union{Bool, Tuple}}
│         (p = Core.getfield(%9, 1))
│   %11 = Core.getfield(%9, 2)::Union{Bool, Tuple}
│         Main.println(p)
│         (@_4 = Base.iterate(%4, %11))
│   %14 = (@_4::Union{Nothing, Tuple{Any, Tuple{Any, Vararg{Any}}}} === nothing)::Bool
│   %15 = Base.not_int(%14)::Bool
└──       goto #4 if not %15
3 ─       goto #2
4 ┄       return nothing
  1. A hack using Combinatorics.multiset_permutations. This approach yields type-stable code but requires duplicating the input.
    Example:
using Combinatorics

function test_combinatorics(n , k)
    x = collect(1:n)
    duplicates = append!(x, x)
    for p in multiset_permutations(duplicates, k)
        println(p)
    end
end

Output of @code_warntype:

julia> @code_warntype test_combinatorics(3, 2)
MethodInstance for test_combinatorics(::Int64, ::Int64)
  from test_combinatorics(n, k) in Main at REPL[2]:1
Arguments
  #self#::Core.Const(test_combinatorics)
  n::Int64
  k::Int64
Locals
  @_4::Union{Nothing, Tuple{Vector{Int64}, Vector{Int64}}}
  duplicates::Vector{Int64}
  x::Vector{Int64}
  p::Vector{Int64}
Body::Nothing
1 ─ %1  = (1:n)::Core.PartialStruct(UnitRange{Int64}, Any[Core.Const(1), Int64])
│         (x = Main.collect(%1))
│         (duplicates = Main.append!(x, x))
│   %4  = Main.multiset_permutations(duplicates, k)::Combinatorics.MultiSetPermutations{Vector{Int64}}
│         (@_4 = Base.iterate(%4))
│   %6  = (@_4 === nothing)::Bool
│   %7  = Base.not_int(%6)::Bool
└──       goto #4 if not %7
2 ┄ %9  = @_4::Tuple{Vector{Int64}, Vector{Int64}}
│         (p = Core.getfield(%9, 1))
│   %11 = Core.getfield(%9, 2)::Vector{Int64}
│         Main.println(p)
│         (@_4 = Base.iterate(%4, %11))
│   %14 = (@_4 === nothing)::Bool
│   %15 = Base.not_int(%14)::Bool
└──       goto #4 if not %15
3 ─       goto #2
4 ┄       return nothing

Could anybody point me to the principled way to do what I want in Julia?

if k can be known in this case k=2:


julia> function test_base_iterators(n)
           Base.Cartesian.@nloops 2 prefix _->1:n begin
               p = Base.Cartesian.@ntuple 2 prefix
               @show p
           end
       end

is type-stable


also this makes a bunch of allocation right?

1 Like

Actually I am interested in the case where k is only known at runtime.

conceptually I don’t think it’s possible to make it type stable:

  1. the cartsian power type is Tuple{blah, blah ...}, with different number, you get different type
  2. how many blah insides Tuple’s parametric type depends on VALUE of k, not type of k

→ not compatible with the definition “types of local variables / outputs can be inferred from types of inputs”

2 Likes

Well that’s too bad. If only multiset_permutations returned a Cartesian power for unique elements.

The concept of “runtime” is somewhat blurred in Julia, because Julia has a runtime compiler.

If k is determined at runtime, but then you run *lots * of computations with a particular k, then it makes sense to encode k in a type, e.g. as Val{k} or the length of an SVector from StaticArrays.jl. Then you can force the compiler to compile specialized code for that k (e.g. with @nloops in an @generated function, which takes a little while but is worth it if the subsequent computations are long.

Otherwise, the generic way to nest arbitrarily many (k) loops with a runtime depth k is to use recursion.

1 Like

e.g. something like:

function _recursive_iterator!(x, n, k, i)
    if i > k
        println(x)
    elseif i == k # unroll 1-loop base case for speed
        for j = 1:n
            x[i] = j
            println(x)
        end
    else
        for j = 1:n
            x[i] = j
            _recursive_iterator!(x, n, k, i+1)
        end
    end
end
function recursive_iterator(n, k)
    x = zeros(Int, k)
    _recursive_iterator!(x, n, k, 1)
end

for example:

julia> recursive_iterator(3,2)
[1, 1]
[1, 2]
[1, 3]
[2, 1]
[2, 2]
[2, 3]
[3, 1]
[3, 2]
[3, 3]

Note that this is type-stable and allocation-free in the inner loops because it re-uses the same vector x over and over, and by unrolling (“recursion coarsening”) a large enough base case (e.g. you could unroll for k==2 and k==3 as well) you can make it quite efficient.

If you want a type-stable container with an arbitrary length, just use a Vector, not a tuple.

2 Likes