CLN: Use public APIs (#198)

* CLN: Use public APIs

* Use compat for temp dir

* Fix linkt

Author: seanpmorgan

tensorflow_addons/image/distance_transform.py

@@ -25,7 +25,7 @@
 _image_ops_so = tf.load_op_library(
     get_path_to_datafile("custom_ops/image/_image_ops.so"))
 
-ops.NotDifferentiable("EuclideanDistanceTransform")
+tf.no_gradient("EuclideanDistanceTransform")
 ops.RegisterShape("EuclideanDistanceTransform")(
     common_shapes.call_cpp_shape_fn)
 

tensorflow_addons/losses/lifted.py

@@ -19,8 +19,6 @@
 from __future__ import print_function
 
 import tensorflow as tf
-
-from tensorflow.python.keras import losses
 from tensorflow_addons.losses import metric_learning
 from tensorflow_addons.utils import keras_utils
 
@@ -106,7 +104,7 @@ def lifted_struct_loss(labels, embeddings, margin=1.0):
 
 
 @keras_utils.register_keras_custom_object
-class LiftedStructLoss(losses.LossFunctionWrapper):
+class LiftedStructLoss(keras_utils.LossFunctionWrapper):
     """Computes the lifted structured loss.
 
     The loss encourages the positive distances (between a pair of embeddings

tensorflow_addons/losses/triplet.py

@@ -18,7 +18,6 @@
 from __future__ import print_function
 
 import tensorflow as tf
-from tensorflow.python.keras import losses
 from tensorflow_addons.losses import metric_learning
 from tensorflow_addons.utils import keras_utils
 
@@ -134,7 +133,7 @@ def triplet_semihard_loss(y_true, y_pred, margin=1.0):
 
 
 @keras_utils.register_keras_custom_object
-class TripletSemiHardLoss(losses.LossFunctionWrapper):
+class TripletSemiHardLoss(keras_utils.LossFunctionWrapper):
     """Computes the triplet loss with semi-hard negative mining.
 
     The loss encourages the positive distances (between a pair of embeddings

tensorflow_addons/optimizers/lazy_adam_test.py

@@ -21,8 +21,6 @@
 import numpy as np
 import tensorflow as tf
 
-from tensorflow.python.eager import context
-from tensorflow.python.ops import variables
 from tensorflow_addons.optimizers import lazy_adam
 from tensorflow_addons.utils import test_utils
 
@@ -83,7 +81,7 @@ def testSparse(self):
                 opt = lazy_adam.LazyAdam()
                 update = opt.apply_gradients(
                     zip([grads0, grads1], [var0, var1]))
-                self.evaluate(variables.global_variables_initializer())
+                self.evaluate(tf.compat.v1.global_variables_initializer())
 
                 # Fetch params to validate initial values
                 self.assertAllClose([1.0, 1.0, 2.0], self.evaluate(var0))
@@ -120,7 +118,7 @@ def testSparseDevicePlacement(self):
                 g_sum = lambda: tf.math.reduce_sum(tf.gather(var, indices))  # pylint: disable=cell-var-from-loop
                 optimizer = lazy_adam.LazyAdam(3.0)
                 minimize_op = optimizer.minimize(g_sum, var_list=[var])
-                self.evaluate(variables.global_variables_initializer())
+                self.evaluate(tf.compat.v1.global_variables_initializer())
                 self.evaluate(minimize_op)
 
     @test_utils.run_deprecated_v1
@@ -143,7 +141,7 @@ def testSparseRepeatedIndices(self):
                 aggregated_update_opt = lazy_adam.LazyAdam()
                 aggregated_update = aggregated_update_opt.apply_gradients(
                     [(grad_aggregated, aggregated_update_var)])
-                self.evaluate(variables.global_variables_initializer())
+                self.evaluate(tf.compat.v1.global_variables_initializer())
                 self.assertAllClose(aggregated_update_var.eval(),
                                     repeated_index_update_var.eval())
                 for _ in range(3):
@@ -182,10 +180,10 @@ def doTestBasic(self, use_callable_params=False):
                     epsilon = epsilon()
 
                 opt = lazy_adam.LazyAdam(learning_rate=learning_rate)
-                if not context.executing_eagerly():
+                if not tf.executing_eagerly():
                     update = opt.apply_gradients(
                         zip([grads0, grads1], [var0, var1]))
-                    self.evaluate(variables.global_variables_initializer())
+                    self.evaluate(tf.compat.v1.global_variables_initializer())
                     # Fetch params to validate initial values
                     self.assertAllClose([1.0, 2.0], self.evaluate(var0))
                     self.assertAllClose([3.0, 4.0], self.evaluate(var1))
@@ -198,7 +196,7 @@ def doTestBasic(self, use_callable_params=False):
                         0.9**(t + 1), self.evaluate(beta_1_power))
                     self.assertAllCloseAccordingToType(
                         0.999**(t + 1), self.evaluate(beta_2_power))
-                    if not context.executing_eagerly():
+                    if not tf.executing_eagerly():
                         self.evaluate(update)
                     else:
                         opt.apply_gradients(
@@ -222,8 +220,7 @@ def testResourceBasic(self):
         self.doTestBasic()
 
     def testBasicCallableParams(self):
-        with context.eager_mode():
-            self.doTestBasic(use_callable_params=True)
+        self.doTestBasic(use_callable_params=True)
 
     @test_utils.run_deprecated_v1
     def testTensorLearningRate(self):
@@ -243,7 +240,7 @@ def testTensorLearningRate(self):
                 opt = lazy_adam.LazyAdam(tf.constant(0.001))
                 update = opt.apply_gradients(
                     zip([grads0, grads1], [var0, var1]))
-                self.evaluate(variables.global_variables_initializer())
+                self.evaluate(tf.compat.v1.global_variables_initializer())
 
                 # Fetch params to validate initial values
                 self.assertAllClose([1.0, 2.0], var0.eval())
@@ -289,7 +286,7 @@ def testSharing(self):
                     zip([grads0, grads1], [var0, var1]))
                 update2 = opt.apply_gradients(
                     zip([grads0, grads1], [var0, var1]))
-                self.evaluate(variables.global_variables_initializer())
+                self.evaluate(tf.compat.v1.global_variables_initializer())
 
                 beta_1_power, beta_2_power = get_beta_accumulators(opt, dtype)
 
@@ -320,16 +317,14 @@ def testSharing(self):
                                                        self.evaluate(var1))
 
     def testSlotsUniqueEager(self):
-        with context.eager_mode():
-            v1 = tf.Variable(1.)
-            v2 = tf.Variable(1.)
-            opt = lazy_adam.LazyAdam(1.)
-            opt.minimize(lambda: v1 + v2, var_list=[v1, v2])
-            # There should be iteration, and two unique slot variables for v1 and v2.
-            self.assertEqual(5, len(set(opt.variables())))
-            self.assertEqual(
-                self.evaluate(opt.variables()[0]),
-                self.evaluate(opt.iterations))
+        v1 = tf.Variable(1.)
+        v2 = tf.Variable(1.)
+        opt = lazy_adam.LazyAdam(1.)
+        opt.minimize(lambda: v1 + v2, var_list=[v1, v2])
+        # There should be iteration, and two unique slot variables for v1 and v2.
+        self.assertEqual(5, len(set(opt.variables())))
+        self.assertEqual(
+            self.evaluate(opt.variables()[0]), self.evaluate(opt.iterations))
 
 
 if __name__ == "__main__":

tensorflow_addons/text/skip_gram_ops_test.py

@@ -22,7 +22,6 @@
 import os
 import tensorflow as tf
 
-from tensorflow.compat.v1 import test
 from tensorflow_addons import text
 from tensorflow_addons.text import skip_gram_ops
 from tensorflow_addons.utils import test_utils
@@ -368,7 +367,8 @@ def test_filter_input_subsample_vocab(self):
 
     @staticmethod
     def _make_text_vocab_freq_file():
-        filepath = os.path.join(test.get_temp_dir(), "vocab_freq.txt")
+        filepath = os.path.join(tf.compat.v1.test.get_temp_dir(),
+                                "vocab_freq.txt")
         with open(filepath, "w") as f:
             writer = csv.writer(f)
             writer.writerows([
@@ -382,7 +382,8 @@ def _make_text_vocab_freq_file():
 
     @staticmethod
     def _make_text_vocab_float_file():
-        filepath = os.path.join(test.get_temp_dir(), "vocab_freq_float.txt")
+        filepath = os.path.join(tf.compat.v1.test.get_temp_dir(),
+                                "vocab_freq_float.txt")
         with open(filepath, "w") as f:
             writer = csv.writer(f)
             writer.writerows([

tensorflow_addons/utils/keras_utils.py

@@ -17,11 +17,12 @@
 from __future__ import division
 from __future__ import print_function
 
+import tensorflow as tf
+
 # TODO: find public API alternative to these
 from tensorflow.python.keras.losses import LossFunctionWrapper  # pylint: disable=unused-import
-from tensorflow.keras.utils import get_custom_objects
 
 
 def register_keras_custom_object(cls):
-    get_custom_objects()[cls.__name__] = cls
+    tf.keras.utils.get_custom_objects()[cls.__name__] = cls
     return cls