CODEWITHSUNDEEP
pythonc++tensorflow
Ayushi Agarwal
Ayushi Agarwal

Reputation: 460

Can't get a new tf.Operation to work in Python shell in Tensorflow

I am trying to add a new integer Matrix Multiplication OP in tensorflow and I am not able to successfully register it as a tf operation so that it can be called as tf.intmatmul in python.

Steps I did : 1) Added a new REGISTER_OP - IntMatMul in the math_ops.cc file.

2) Added a new kernel implementation for this OP in the core/kernels path - int_matmul_op.cc and a corresponding header file - int_matmul_op.h

3) I added the dependency of the OP in the core/kernels/BUILD file. This will add the kernel linking for this OP.

4) Added the definition for this OP (as 'intmatmul') in the Python Wrapper file i.e python/ops/math_ops.py - This file calls the gen_math_ops.int_mat_mul

5) re-built from source using Bazel and re-installed Tensorflow using the pip package.

However when I try to use this OP as tf.intmatmul, I get an error saying the module is not defined. I am not sure now what I am missing here. Is there any linking that is missing? Do I need to add any OP linking in the core/BUILD file as well?

Upvotes: 1

Views: 196

Answers (1)

javidcf
javidcf

Reputation: 59731

UPDATE:

So this turned out to be more complicated than expected. These are the things that had to be taken into account:

  • Apparently, in order for a function to be exposed as public API (that is, at tf. level), its name must be listed at the beginning of the module in its docstring preceded with @@. Take a look for example to math_ops.py.
  • The Kernel definition has to be absolutely correct in order to reflect an operation as an operation, even if you can still access it from the internal module (e.g. doing from tensorflow.python.ops import math_ops).

--

As the docs indicate, the name of the operation, that must be registered with a CamelCase identifier in C++, is "translated" into snake_case in Python. Try tf.int_mat_mul instead.

As a side note, that tutorial provides additional guidance to implement custom operations without needing to recompile TensorFlow from source, loading it from a custom library instead.

Upvotes: 2

Related Questions