Google I/O is a wrap! Catch up on TensorFlow sessions View sessions


View source on GitHub

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

This function is a wrapper around tf.compat.v1.py_func and improve it with kwargs and output_shapes. Further it changed some argument names.

Given a python function func, which takes numpy arrays as its inputs 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)
inp = tf.compat.v1.placeholder(tf.float32)
y = tf.compat.v1.py_func(my_func, [inp], 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():).

func A Python function, which accepts a list of NumPy ndarray objects having element types that match the corresponding tf.Tensor objects in inp, and returns a list of ndarray objects (or a single ndarray) having element types that match the corresponding values in Tout.
args A list of Tensor objects.
kwargs A dict with Tensor objects as values.
output_types A nested structure of tensorflow data types or a single tensorflow data type if there is only one, indicating what func returns.
output_shapes Same as output_types, except the types are replaces with shapes (optional).
stateful (Boolean.) If True, the function should be considered stateful. If a function is stateless, when given the same input it will return the same output and have no observable side effects. Optimizations such as common subexpression elimination are only performed on stateless operations.
name A name for the operation (optional).

Tensorflow op that wraps the input python function.