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])
Args |
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.
|
Returns |
A new function that applies pyfunc to every element
of its input along axis 0 (the batch axis).
|