tf.numpy_function  |  TensorFlow v2.16.1 (original) (raw)

tf.numpy_function

Stay organized with collections Save and categorize content based on your preferences.

Wraps a python function and uses it as a TensorFlow op.

View aliases

Compat aliases for migration

SeeMigration guide for more details.

tf.compat.v1.numpy_function

tf.numpy_function(
    func=None, inp=None, Tout=None, stateful=True, name=None
)

Used in the notebooks

Used in the tutorials
Playing CartPole with the Actor-Critic method

Given a python function func wrap this function as an operation in atf.function. func must take numpy arrays as its arguments and return numpy arrays as its outputs.

There are two ways to use tf.numpy_function.

As a decorator

When using tf.numpy_function as a decorator:

@tf.numpy_function(Tout=tf.float32) def my_numpy_func(x): # x will be a numpy array with the contents of the input to the # tf.function print(f'executing eagerly, {x=}') return np.sinh(x)

The function runs eagerly:

my_numpy_func(1.0).numpy() executing eagerly, x=1.0 1.17520

The behavior doesn't change inside a tf.function:

@tf.function(input_signature=[tf.TensorSpec(None, tf.float32)]) def tf_function(input): y = tf.numpy_function(my_numpy_func, [input], tf.float32) return y tf_function(tf.constant(1.)).numpy() executing eagerly, x=array(1.) 1.17520

Inplace

This form can be useful if you don't control the function's source, but it is harder to read.

Here is the same function with no decorator:

def my_func(x): # x will be a numpy array with the contents of the input to the # tf.function print(f'executing eagerly, {x=}') return np.sinh(x)

To run tf.numpy_function in-place, pass the function, its inputs, and the output type in a single call to tf.numpy_function:

tf.numpy_function(my_func, [tf.constant(1.0)], tf.float32) executing eagerly, x=array(1.) 1.17520

More info

Comparison to tf.py_function:tf.py_function and tf.numpy_function are very similar, except thattf.numpy_function takes numpy arrays, and not tf.Tensors. If you want the function to contain tf.Tensors, and have any TensorFlow operations executed in the function be differentiable, please use tf.py_function.

Args
func A Python function, which accepts numpy.ndarray objects as arguments and returns a list of numpy.ndarray objects (or a singlenumpy.ndarray). This function must accept as many arguments as there are tensors in inp, and these argument types will match the correspondingtf.Tensor objects in inp. The returns numpy.ndarrays must match the number and types defined Tout. Important Note: Input and outputnumpy.ndarrays of func are not guaranteed to be copies. In some cases their underlying memory will be shared with the corresponding TensorFlow tensors. In-place modification or storing func input or return values in python datastructures without explicit (np.)copy can have non-deterministic consequences.
inp A list of tf.Tensor objects.
Tout A list or tuple of tensorflow data types or a single tensorflow data type if there is only one, indicating what func returns.
stateful (Boolean.) Setting this argument to False tells the runtime to treat the function as stateless, which enables certain optimizations. A function is stateless when given the same input it will return the same output and have no side effects; its only purpose is to have a return value. The behavior for a stateful function with the stateful argument False is undefined. In particular, caution should be taken when mutating the input arguments as this is a stateful operation.
name (Optional) A name for the operation.
Returns
If func is None this returns a decorator that will ensure the decorated function will always run with eager execution even if called from a tf.function/tf.Graph. If used func is not None this executes func with eager execution and returns the result: A single or list of tf.Tensor which funccomputes.