[Feature] qjit-compatible implementation of `itertools.product`

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.

Start by reading the AutoGraph conversion handling for itertools.product and the existing catalyst.range implementation, then trace how supported constructs lower through the MLIR layer. Done means catalyst.product supports the described qjit loop, including repeat and dynamic arguments, with coverage for the requested semantics and compilation behavior.