Migrate api style

Author: fsx950223

tensorflow_addons/activations/BUILD

@@ -27,7 +27,7 @@
 
 @keras_utils.register_keras_custom_object
 @tf.function
-def tanhshrink(features, name="Tanhshrink"):
+def tanhshrink(features, name=None):
     """Applies the element-wise function: x - tanh(x)
 
     Args:
@@ -36,13 +36,11 @@ def tanhshrink(features, name="Tanhshrink"):
     Returns:
         A `Tensor`. Has the same type as `features`.
     """
-    with tf.name_scope(name) as name:
+    with tf.name_scope(name or "tanhshrink") as name:
         features = tf.convert_to_tensor(features, name="features")
-        if features.dtype.is_integer:
-            features = tf.cast(features, tf.float32)
-        return _activation_ops_so.tanhshrink(features, name=name)
+        return _activation_ops_so.addons_tanhshrink(features, name=name)
 
 
-@tf.RegisterGradient("Tanhshrink")
+@tf.RegisterGradient("Addons>Tanhshrink")
 def _tanhshrink_grad(op, grad):
-    return _activation_ops_so.tanhshrink_grad(grad, op.inputs[0])
+    return _activation_ops_so.addons_tanhshrink_grad(grad, op.inputs[0])

tensorflow_addons/activations/tanhshrink.py

@@ -96,6 +96,6 @@ cc_binary(
     ] + if_cuda_is_configured([
         ":gelu_op_gpu",
         ":hardshrink_op_gpu",
-        ":tanhshrink_op_gpu"
+        ":tanhshrink_op_gpu",
     ]),
 )

tensorflow_addons/custom_ops/activations/BUILD

@@ -21,16 +21,18 @@ limitations under the License.
 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
 
 namespace tensorflow {
+namespace addons {
 
 using CPUDevice = Eigen::ThreadPoolDevice;
 
-#define REGISTER_TANHSHRINK_KERNELS(type)                                  \
-  REGISTER_KERNEL_BUILDER(                                                 \
-      Name("Tanhshrink").Device(DEVICE_CPU).TypeConstraint<type>("T"),     \
-      TanhshrinkOp<CPUDevice, type>);                                      \
-  REGISTER_KERNEL_BUILDER(                                                 \
-      Name("TanhshrinkGrad").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
-      TanhshrinkGradOp<CPUDevice, type>);
+#define REGISTER_TANHSHRINK_KERNELS(type)                                     \
+  REGISTER_KERNEL_BUILDER(                                                    \
+      Name("Addons>Tanhshrink").Device(DEVICE_CPU).TypeConstraint<type>("T"), \
+      TanhshrinkOp<CPUDevice, type>);                                         \
+  REGISTER_KERNEL_BUILDER(Name("Addons>TanhshrinkGrad")                       \
+                              .Device(DEVICE_CPU)                             \
+                              .TypeConstraint<type>("T"),                     \
+                          TanhshrinkGradOp<CPUDevice, type>);
 
 TF_CALL_GPU_NUMBER_TYPES(REGISTER_TANHSHRINK_KERNELS);
 #undef REGISTER_TANHSHRINK_KERNELS
@@ -60,17 +62,19 @@ TF_CALL_GPU_NUMBER_TYPES(DECLARE_GPU_SPEC);
 }  // namespace functor
 
 // Registration of the GPU implementations.
-#define REGISTER_TANHSHRINK_GPU_KERNELS(type)                              \
-  REGISTER_KERNEL_BUILDER(                                                 \
-      Name("Tanhshrink").Device(DEVICE_GPU).TypeConstraint<type>("T"),     \
-      TanhshrinkOp<GPUDevice, type>);                                      \
-  REGISTER_KERNEL_BUILDER(                                                 \
-      Name("TanhshrinkGrad").Device(DEVICE_GPU).TypeConstraint<type>("T"), \
-      TanhshrinkGradOp<GPUDevice, type>);
+#define REGISTER_TANHSHRINK_GPU_KERNELS(type)                                 \
+  REGISTER_KERNEL_BUILDER(                                                    \
+      Name("Addons>Tanhshrink").Device(DEVICE_GPU).TypeConstraint<type>("T"), \
+      TanhshrinkOp<GPUDevice, type>);                                         \
+  REGISTER_KERNEL_BUILDER(Name("Addons>TanhshrinkGrad")                       \
+                              .Device(DEVICE_GPU)                             \
+                              .TypeConstraint<type>("T"),                     \
+                          TanhshrinkGradOp<GPUDevice, type>);
 
 TF_CALL_GPU_NUMBER_TYPES(REGISTER_TANHSHRINK_GPU_KERNELS);
 #undef REGISTER_TANHSHRINK_GPU_KERNELS
 
 #endif  // GOOGLE_CUDA
 
+}  // namespace addons
 }  // namespace tensorflow

tensorflow_addons/custom_ops/activations/cc/kernels/tanhshrink_op.cc

@@ -13,8 +13,8 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#ifndef TENSORFLOW_ADDONS_TANHSHRINK_OP_H_
-#define TENSORFLOW_ADDONS_TANHSHRINK_OP_H_
+#ifndef TENSORFLOW_ADDONS_ACTIVATIONS_KERNELS_TANHSHRINK_OP_H_
+#define TENSORFLOW_ADDONS_ACTIVATIONS_KERNELS_TANHSHRINK_OP_H_
 
 #define EIGEN_USE_THREADS
 
@@ -23,6 +23,7 @@ limitations under the License.
 #include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
 
 namespace tensorflow {
+namespace addons {
 namespace functor {
 
 template <typename Device, typename T>
@@ -88,7 +89,8 @@ void TanhshrinkGradOp<Device, T>::OperateNoTemplate(OpKernelContext* context,
   functor(context->eigen_device<Device>(), g.flat<T>(), a.flat<T>(),
           output->flat<T>());
 }
-}
+}  // namespace addons
+}  // namespace tensorflow
 
 #undef EIGEN_USE_THREADS
 

tensorflow_addons/custom_ops/activations/cc/kernels/tanhshrink_op.h

@@ -22,6 +22,7 @@ limitations under the License.
 #include "third_party/eigen3/Eigen/Core"
 
 namespace tensorflow {
+namespace addons {
 
 using GPUDevice = Eigen::GpuDevice;
 
@@ -31,6 +32,7 @@ using GPUDevice = Eigen::GpuDevice;
 
 TF_CALL_GPU_NUMBER_TYPES(DEFINE_GPU_KERNELS);
 
+}  // namespace addons
 }  // namespace tensorflow
 
 #endif  // GOOGLE_CUDA

tensorflow_addons/custom_ops/activations/cc/kernels/tanhshrink_op_gpu.cu.cc

@@ -18,18 +18,20 @@ limitations under the License.
 #include "tensorflow/core/framework/shape_inference.h"
 
 namespace tensorflow {
+namespace addons {
 
-REGISTER_OP("Tanhshrink")
+REGISTER_OP("Addons>Tanhshrink")
     .Input("features: T")
     .Output("activations: T")
     .Attr("T: {half, float, double}")
     .SetShapeFn(shape_inference::UnchangedShape);
 
-REGISTER_OP("TanhshrinkGrad")
+REGISTER_OP("Addons>TanhshrinkGrad")
     .Input("gradients: T")
     .Input("features: T")
     .Output("backprops: T")
     .Attr("T: {half, float, double}")
     .SetShapeFn(shape_inference::MergeBothInputsShapeFn);
 
+}  // namespace addons
 }  // namespace tensorflow