def _product()

in src/beanmachine/ppl/compiler/support.py [0:0]

• 11 lines of code
• 18 McCabe index (conditional complexity)

    def _product(self, f: Callable, *nodes: bn.BMGNode) -> SetOfTensors:
        # * We have a sequence of nodes n[0], n[1], ... n[k-1].
        #
        # * Each of those nodes has possible values t[x][0], t[x][1] ...
        #   for x from 0 to k-1.
        #
        # * We have a function f which takes k tensors and returns a tensor.
        #
        # * We wish to compute the set:
        #
        # {
        #   f(t[0][0], t[1][0], ... t[k-1][0]),
        #   f(t[0][1], t[1][0], ... t[k-1][0]),
        #   ...
        # }
        #
        # That is, the Cartesian product of all possible combinations of
        # values of each node, with each element of the product run through
        # function f.
        #
        # However, we have some problems:
        #
        # * The support of a node might be infinite.
        # * The support of a node might be unknown.
        # * The support of a node might be finite but large.
        # * The size of the product might be finite but large.
        #
        # In those cases we want to return special values Infinite, Unknown
        # or TooBig rather than wasting time and memory to compute the set.
        #
        # ####
        #
        # First thing to do is determine the *approximate* size of the
        # Cartesian product of possible values.
        #
        # TODO: This approximation is an over-estimate; for instance, when
        # multiplying {0 or 1} by { n elements} we assume that the resulting
        # set has up to 2n elements, when in fact it only has n or n+1 elements.
        # Would it be simpler and more accurate to instead make a loop, accumulate
        # the result into a mutable deduplicating set, and if the set ever gets
        # too big, bail out then?

        size = functools.reduce(
            lambda acc, node: _set_product_approximate_size(acc, self[node]), nodes, 1.0
        )

        if size == positive_infinity:
            return Infinite
        if isnan(size):
            return Unknown
        if size >= _limit:
            return TooBig

        # If we've made it here then every node had a small support
        # and the product of the approximate sizes was small too.
        # We can just take the Cartesian product and call f.

        return SetOfTensors(f(c) for c in itertools.product(*(self[n] for n in nodes)))