Skip to content

vectorize

vectorize #

vectorize — lift a scalar function so it broadcasts over Array arguments.

@vectorize
def add(a, b):
    return a + b


add(1, 2)  # 3 — scalar passthrough
add(Array([1, 2, 3]), 10)  # Array([11, 12, 13])

A length-1 sequence or a scalar is repeated to match the longest sequence; two sequences of different non-unit lengths are incompatible.

vectorize(function) #

Wrap function so Array arguments broadcast element-wise.

With no sequence argument the call is a plain scalar call; otherwise every argument broadcasts to the common length and the results collect into an Array.

Source code in capturegraph-lib/capturegraph/types/containers/vectorize.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
def vectorize(function: Callable) -> Callable:
    """Wrap ``function`` so ``Array`` arguments broadcast element-wise.

    With no sequence argument the call is a plain scalar call; otherwise every
    argument broadcasts to the common length and the results collect into an
    ``Array``.
    """

    @functools.wraps(function)
    def wrapper(*args: object, **kwargs: object) -> object:
        if not any(_is_sequence(value) for value in (*args, *kwargs.values())):
            return function(*args, **kwargs)
        length = _broadcast_length(*args, *kwargs.values())
        if length == 0:
            return Array()
        rows_args = [_broadcast_to(arg, length) for arg in args]
        rows_kwargs = {key: _broadcast_to(value, length) for key, value in kwargs.items()}
        return Array(
            function(
                *(column[row] for column in rows_args),
                **{key: column[row] for key, column in rows_kwargs.items()},
            )
            for row in range(length)
        )

    return wrapper