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])

pyfuncCallable of a single tensor argument.
excludedOptional set of integers representing positional arguments for which the function will not be vectorized. These will be passed directly to pyfunc unmodified.
signatureOptional 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).