tf.numpy_function

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

Used in the notebooks

Used in the tutorials

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 mu