[Feature] qjit-compatible implementation of itertools.product

[josh146]

It's quite common to see workflows that make use of itertools related functionality, in particular

• itertools.product,
• itertools.combinations, and
• itertools.permutations.

Here, I focus only on product as this is part of the workflow I am using.

For example, consider:

@qml.qjit(autograph=True)
def fn(x):

    for i, j in itertools.product(range(2), repeat=2):
        x += i + j - i * j * jnp.sin(j)

    return x

Here, product has the following rough semantics:

def p(*iterables, repeat=1):
    result = [[]]

    if repeat < 0:
        raise ValueError('repeat argument cannot be negative')

    for l in range(repeat):
        for pool in map(tuple, iterables):
            _temp = []

            for i in result:
                for j in pool:
                    _temp.append(i + [j])

            result = _temp

    for prod in result:
        yield tuple(prod)

(note this is not exactly the implementation, as the actual implementation does not generate intermediate lists in memory).

This will fail to compile, since Autograph is not able to recognize and convert itertools.product (even though the arguments are all compile-time constant).

It is also non-trivial to re-write as pure catalyst loops using range in complex cases.

I propose adding our own catalyst.product, similar to our existing range, that AutoGraph can use as the conversion target. This would then lower to a fast implementation of product (likely as nested for loops) at the MLIR layer. An advantage here is that this would allow product to work even with dynamic arguments.

[josh146]

Longer term, we could also do the same thing for itertools.combinations and itertools.permutations.