Tensorflow: Writing an Op in Python

Question

I would like to write an Op in Python. This tutorial only explains how to do it in c++ with a Python wrapper. https://www.tensorflow.org/versions/master/how_tos/adding_an_op/index.html#adding-a-new-op

How can I write it completely in Python?

Answer 1

You can use tf.py_func(func, inp, Tout).

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

Given a python function func, which takes numpy arrays as its inputs and returns numpy arrays as its outputs.

---

Your python function needs to have:

• numpy arrays as inputs, fed from the graph with the argument inp
• numpy arrays as outputs, you need to specify their types to TensorFlow in the argument Tout

Inside the function, you can do whatever you like, if conditions of for loops, anything that is not possible in TensorFlow.

---

However, the operation will be executed on CPU so it may be slower than the equivalent TensorFlow op in GPU.