tf.numpy_function

View source on GitHub

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

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

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

The following example creates a TensorFlow graph with np.sinh() as an operation in the graph:

def my_numpy_func(x):
  # x will be a numpy array with the contents of the input to the
  # tf.function
  return np.sinh(x)
@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 * y
tf_function(tf.constant(1.))
<tf.Tensor: shape=(), dtype=float32, numpy=1.3810978>

Comparison to tf.py_function: tf.py_function and tf.numpy_function are very similar, except that tf.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.

• The body of the function (i.e. func) will not be serialized in a tf.SavedModel. Therefore, you should not use this function if you need to serialize your model and restore it in a different environment.

• The operation must run in the same address space as the Python program that calls tf.numpy_function(). If you are using distributed TensorFlow, you must run a tf.distribute.Server in the same process as the program that calls tf.numpy_function you must pin the created operation to a device in that server (e.g. using with tf.device():).

• Since the function takes numpy arrays, you cannot take gradients through a numpy_function. If you require something that is differentiable, please consider using tf.py_function.

• The resulting function is assumed stateful and will never be optimized.