def vectorize_stochastic(orig):
    if orig.name == 'idxs_map' and orig.pos_args[1]._obj in stoch:
        # -- this is an idxs_map of a random draw of distribution `dist`
        idxs = orig.pos_args[0]
        dist = orig.pos_args[1]._obj
        def foo(arg):
            # -- each argument is an idxs, vals pair
            assert arg.name == 'pos_args'
            assert len(arg.pos_args) == 2
            arg_vals = arg.pos_args[1]

            # XXX: write a pattern-substitution rule for this case
            if arg_vals.name == 'idxs_take':
                if arg_vals.arg['vals'].name == 'asarray':
                    if arg_vals.arg['vals'].inputs()[0].name == 'repeat':
                        # -- draws are iid, so forget about
                        #    repeating the distribution parameters
                        repeated_thing = arg_vals.arg['vals'].inputs()[0].inputs()[1]
                        return repeated_thing
            if arg.pos_args[0] is idxs:
                return arg_vals
            else:
                # -- arg.pos_args[0] is a superset of idxs
                #    TODO: slice out correct elements using
                #    idxs_take, but more importantly - test this case.
                raise NotImplementedError()
        new_pos_args = [foo(arg) for arg in orig.pos_args[2:]]
        new_named_args = [[aname, foo(arg)]
                for aname, arg in orig.named_args]
        vnode = Apply(dist, new_pos_args, new_named_args, o_len=None)
        n_times = scope.len(idxs)
        if dict(vnode.named_args) in 'size':
            raise NotImplementedError('random node already has size')
        vnode.named_args.append(['size', n_times])
        return vnode
    else:
        return orig
dataset/ETHPy150Open hyperopt/hyperopt/hyperopt/vectorize.py/vectorize_stochastic/swapped operands of `in`