tf.keras.ops.vectorize

Turn a function into a vectorized function.

tf.keras.ops.vectorize(
    pyfunc, *, excluded=None, signature=None
)

Example:

def myfunc(a, b):
    return a + b

vfunc = np.vectorize(myfunc)
y = vfunc([1, 2, 3, 4], 2)  # Returns Tensor([3, 4, 5, 6])

pyfunc Callable of a single tensor argument.
excluded Optional set of integers representing positional arguments for which the function will not be vectorized. These will be passed directly to pyfunc unmodified.
signature Optional generalized universal function signature, e.g., "(m,n),(n)->(m)" for vectorized matrix-vector multiplication. If provided, pyfunc will be called with (and expected to return) arrays with shapes given by the size of corresponding core dimensions. By default, pyfunc is assumed to take scalars tensors as input and output.

A new function that applies pyfunc to every element of its input along axis 0 (the batch axis).