pytensor.tensor.functional.vectorize(func: Callable, signature: Optional[str] = None) → Callable

Create a vectorized version of a python function that takes TensorVariables as inputs and outputs.

Similar to numpy.vectorize. See respective docstrings for more details.

Parameters:

• func (Callable) – Function that creates the desired outputs from TensorVariable inputs with the core dimensions.

• signature (str, optional) – Generalized universal function signature, e.g., (m,n),(n)->(m) for vectorized matrix-vector multiplication. If not provided, it is assumed all inputs have scalar core dimensions. Unlike numpy, the outputs can have arbitrary shapes when the signature is not provided.

Returns:

vectorized_func – Callable that takes TensorVariables with arbitrarily batched dimensions on the left and returns variables whose graphs correspond to the vectorized expressions of func.

Return type:

Callable

Notes

Unlike numpy.vectorize, the equality of core dimensions implied by the signature is not explicitly asserted.

To vectorize an existing graph, use pytensor.graph.replace.vectorize_graph instead.

Examples

import pytensor
import pytensor.tensor as pt

def func(x):
    return pt.exp(x) / pt.sum(pt.exp(x))


vec_func = pt.vectorize(func, signature="(a)->(a)")

x = pt.matrix("x")
y = vec_func(x)

fn = pytensor.function([x], y)
fn([[0, 1, 2], [2, 1, 0]])
# array([[0.09003057, 0.24472847, 0.66524096],
#        [0.66524096, 0.24472847, 0.09003057]])

import pytensor
import pytensor.tensor as pt


def func(x):
    return x[0], x[-1]


vec_func = pt.vectorize(func, signature="(a)->(),()")

x = pt.matrix("x")
y1, y2 = vec_func(x)

fn = pytensor.function([x], [y1, y2])
fn([[-10, 0, 10], [-11, 0, 11]])
# [array([-10., -11.]), array([10., 11.])]