tensorflow Creating a custom operation with tf.py_func (CPU only) Basic example


The tf.py_func(func, inp, Tout) operator creates a TensorFlow operation that calls a Python function, func on a list of tensors inp.

See the documentation for tf.py_func(func, inp, Tout).

Warning: The tf.py_func() operation will only run on CPU. If you are using distributed TensorFlow, the tf.py_func() operation must be placed on a CPU device in the same process as the client.

def func(x):
    return 2*x

x = tf.constant(1.)
res = tf.py_func(func, [x], [tf.float32])
# res is a list of length 1