Tune in to the first Women in ML Symposium this Tuesday, October 19 at 9am PST Register now


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

Migrate to TF2

This name was deprecated and removed in TF2, but tf.numpy_function is a near-exact replacement, just drop the stateful argument (all tf.numpy_function calls are considered stateful). It is compatible with eager execution and tf.function.

tf.py_function is a close but not an exact replacement, passing TensorFlow tensors to the wrapped function instead of NumPy arrays, which provides gradients and can take advantage of accelerators.


def fn_using_numpy(x):
  x[0] = 0.
  return x
tf.compat.v1.py_func(fn_using_numpy, inp=[tf.constant([1., 2.])],
    Tout=tf.float32, stateful=False)
<tf.Tensor: shape=(2,), dtype=float32, numpy=array([0., 2.], dtype=float32)>


tf.numpy_function(fn_using_numpy, inp=[tf.constant([1., 2.])],
<tf.Tensor: shape=(2,), dtype=float32, numpy=array([0., 2.], dtype=float32)>


Used in the notebooks

Used in the tutorials

Given a python function func, which takes numpy arrays as its arguments and returns numpy arrays as its outputs, wrap this function as an operation in a TensorFlow graph. The following snippet constructs a simple TensorFlow graph that invokes the np.sinh() NumPy function as a operation in the graph:

def my_func(x):
  # x will be a numpy array with the contents of the placeholder below
  return np.sinh(x)
input = tf.compat.v1.placeholder(tf.float32)
y = tf.compat.v1.py_func(my_func, [input], tf.float32)
  • The body of the function (i.e. func) will not be serialized in a GraphDef. 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.compat.v1.py_func(). If you are using distributed TensorFlow, you must run a tf.distribute.Server in the same process as the program that calls tf.compat.v1.py_func() and you must pin the created operation to a device in that server (e.g. using with tf.device():).


  import tensorflow as tf
  import numpy as np

  def make_synthetic_data(i):
      return np.cast[np.uint8](i) * np.ones([20,256,256,3],
              dtype=np.float32) / 10.

  def preprocess_fn(i):
      ones = tf.py_function(make_synthetic_data,[i],tf.float32)
      ones.set_shape(tf.TensorShape([None, None, None, None]))
      ones = tf.image.resize(ones, [224,224])
      return ones

  ds = tf.data.Dataset.range(10)
  ds = ds.map(preprocess_fn)

Args: func: A Python function, which accepts ndarray objects as arguments and returns a list of ndarray objects (or a single ndarray). This function must accept as many arguments as there are tensors in inp, and these argument types will match the corresponding tf.Tensor objects in inp. The returns ndarrays must match the number and types defined Tout. Important Note: Input and output numpy ndarrays of func are not guaranteed to be copies. In some cases their underlying memory will be shared with the corresponding Tens