tf.nn.sigmoid implementation in TensorFlow

Question

Could you please point me to the actual c++ implementation of the tf.nn.sigmoid function in TensorFlow (2.0) ? Just for educational purposes.

Thanks.

Answer 1

Tensorflow source can be little tricky to navigate. Here's some ideas I use to find implementation files.

You can see that tf.nn.sigmoid is defined in python here

tensorflow/python/ops/math_ops.py.

If you look at this function definition, you can see it returns a sigmoid function defined in gen_math_ops. It took me awhile to figure out, but if you go searching for gen_math_ops in tensorflow/python/ops you won't find it. Anything with gen_ in front is the name tensorflow gives an op registered by C++ code.

What we really want in the kernel implementation. These can be found in tensorflow/core/kernels. A quick ctrl+f for "sigmoid" points us to cwise_op_sigmoid.cc. This doesn't contain the implementation but points us to a header file cwise_ops_common.h. This doesn't contain the implementation either, but points us to cwise_ops.h. A ctrl+f for "sigmoid" in this file and we find the "implementation" on line 877. You can see this is a functor that wraps the Eigen operation Eigen::internal::scalar_logistic_op. Here are the docs for that op. If you're curious how that op is implemented, I'd download the source. You can find it here.