FIX: Use correct TF imports (#55)

* FIX: Use correct TF imports

Author: seanpmorgan

tensorflow_addons/custom_ops/image/python/transform.py

@@ -19,20 +19,14 @@
 
 import tensorflow as tf
 from tensorflow.python.framework import common_shapes
-from tensorflow.python.framework import constant_op
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import load_library
 from tensorflow.python.framework import ops
-from tensorflow.python.ops import array_ops
-from tensorflow.python.ops import linalg_ops
-from tensorflow.python.ops import math_ops
 from tensorflow.python.platform import resource_loader
 
-_image_ops_so = load_library.load_op_library(
+_image_ops_so = tf.load_op_library(
     resource_loader.get_path_to_datafile("_image_ops.so"))
 
-_IMAGE_DTYPES = set([dtypes.uint8, dtypes.int32, dtypes.int64, dtypes.float16,
-                     dtypes.float32, dtypes.float64])
+_IMAGE_DTYPES = set([tf.dtypes.uint8, tf.dtypes.int32, tf.dtypes.int64,
+                     tf.dtypes.float16, tf.dtypes.float32, tf.dtypes.float64])
 
 ops.RegisterShape("ImageProjectiveTransform")(common_shapes.call_cpp_shape_fn)
 
@@ -78,7 +72,7 @@ def transform(images,
         image_or_images = ops.convert_to_tensor(images, name="images")
         transform_or_transforms = ops.convert_to_tensor(transforms,
                                                         name="transforms",
-                                                        dtype=dtypes.float32)
+                                                        dtype=tf.dtypes.float32)
         if image_or_images.dtype.base_dtype not in _IMAGE_DTYPES:
             raise TypeError("Invalid dtype %s." % image_or_images.dtype)
         elif image_or_images.get_shape().ndims is None:
@@ -93,10 +87,10 @@ def transform(images,
             raise TypeError("Images should have rank between 2 and 4.")
 
         if output_shape is None:
-            output_shape = array_ops.shape(images)[1:3]
+            output_shape = tf.shape(images)[1:3]
 
         output_shape = ops.convert_to_tensor(output_shape,
-                                             dtypes.int32,
+                                             tf.dtypes.int32,
                                              name="output_shape")
 
         if not output_shape.get_shape().is_compatible_with([2]):
@@ -147,8 +141,8 @@ def compose_transforms(*transforms):
         composed = flat_transforms_to_matrices(transforms[0])
         for tr in transforms[1:]:
             # Multiply batches of matrices.
-            composed = math_ops.matmul(composed,
-                                       flat_transforms_to_matrices(tr))
+            composed = tf.matmul(composed,
+                                 flat_transforms_to_matrices(tr))
         return matrices_to_flat_transforms(composed)
 
 
@@ -177,15 +171,15 @@ def flat_transforms_to_matrices(transforms):
             raise ValueError("Transforms should be 1D or 2D, got: %s" %
                              transforms)
         # Make the transform(s) 2D in case the input is a single transform.
-        transforms = array_ops.reshape(transforms,
-                                       constant_op.constant([-1, 8]))
-        num_transforms = array_ops.shape(transforms)[0]
+        transforms = tf.reshape(transforms,
+                                tf.constant([-1, 8]))
+        num_transforms = tf.shape(transforms)[0]
         # Add a column of ones for the implicit last entry in the matrix.
-        return array_ops.reshape(
-            array_ops.concat(
-                [transforms, array_ops.ones([num_transforms, 1])],
+        return tf.reshape(
+            tf.concat(
+                [transforms, tf.ones([num_transforms, 1])],
                 axis=1),
-            constant_op.constant([-1, 3, 3]))
+            tf.constant([-1, 3, 3]))
 
 
 def matrices_to_flat_transforms(transform_matrices):
@@ -215,8 +209,8 @@ def matrices_to_flat_transforms(transform_matrices):
             raise ValueError("Matrices should be 2D or 3D, got: %s" %
                              transform_matrices)
         # Flatten each matrix.
-        transforms = array_ops.reshape(transform_matrices,
-                                       constant_op.constant([-1, 9]))
+        transforms = tf.reshape(transform_matrices,
+                                tf.constant([-1, 9]))
         # Divide each matrix by the last entry (normally 1).
         transforms /= transforms[:, 8:9]
         return transforms[:, :8]
@@ -243,28 +237,28 @@ def angles_to_projective_transforms(angles,
     with ops.name_scope(name, "angles_to_projective_transforms"):
         angle_or_angles = ops.convert_to_tensor(angles,
                                                 name="angles",
-                                                dtype=dtypes.float32)
+                                                dtype=tf.dtypes.float32)
         if len(angle_or_angles.get_shape()) == 0:
             angles = angle_or_angles[None]
         elif len(angle_or_angles.get_shape()) == 1:
             angles = angle_or_angles
         else:
             raise TypeError("Angles should have rank 0 or 1.")
         x_offset = ((image_width - 1) -
-                    (math_ops.cos(angles) * (image_width - 1) -
-                     math_ops.sin(angles) * (image_height - 1))) / 2.0
+                    (tf.math.cos(angles) * (image_width - 1) -
+                     tf.math.sin(angles) * (image_height - 1))) / 2.0
         y_offset = ((image_height - 1) -
-                    (math_ops.sin(angles) * (image_width - 1) +
-                     math_ops.cos(angles) * (image_height - 1))) / 2.0
-        num_angles = array_ops.shape(angles)[0]
-        return array_ops.concat(
-            values=[math_ops.cos(angles)[:, None],
-                    -math_ops.sin(angles)[:, None],
+                    (tf.math.sin(angles) * (image_width - 1) +
+                     tf.math.cos(angles) * (image_height - 1))) / 2.0
+        num_angles = tf.shape(angles)[0]
+        return tf.concat(
+            values=[tf.math.cos(angles)[:, None],
+                    -tf.math.sin(angles)[:, None],
                     x_offset[:, None],
-                    math_ops.sin(angles)[:, None],
-                    math_ops.cos(angles)[:, None],
+                    tf.math.sin(angles)[:, None],
+                    tf.math.cos(angles)[:, None],
                     y_offset[:, None],
-                    array_ops.zeros((num_angles, 2), dtypes.float32),],
+                    tf.zeros((num_angles, 2), tf.dtypes.float32),],
             axis=1)
 
 
@@ -278,7 +272,7 @@ def _image_projective_transform_grad(op, grad):
     image_or_images = ops.convert_to_tensor(images, name="images")
     transform_or_transforms = ops.convert_to_tensor(transforms,
                                                     name="transforms",
-                                                    dtype=dtypes.float32)
+                                                    dtype=tf.dtypes.float32)
 
     if image_or_images.dtype.base_dtype not in _IMAGE_DTYPES:
         raise TypeError("Invalid dtype %s." % image_or_images.dtype)
@@ -291,11 +285,11 @@ def _image_projective_transform_grad(op, grad):
 
     # Invert transformations
     transforms = flat_transforms_to_matrices(transforms=transforms)
-    inverse = linalg_ops.matrix_inverse(transforms)
+    inverse = tf.linalg.inv(transforms)
     transforms = matrices_to_flat_transforms(inverse)
     output = _image_ops_so.image_projective_transform(
         images=grad,
         transforms=transforms,
-        output_shape=array_ops.shape(image_or_images)[1:3],
+        output_shape=tf.shape(image_or_images)[1:3],
         interpolation=interpolation)
     return [output, None, None]

tensorflow_addons/custom_ops/image/python/transform_test.py

@@ -19,34 +19,31 @@
 from __future__ import print_function
 
 import numpy as np
+import tensorflow as tf
 
-from tensorflow.python.framework import constant_op
-from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import test_util as tf_test_util
 from tensorflow.python.ops import gradient_checker
-from tensorflow.python.ops import random_ops
-from tensorflow.python.platform import test
 from tensorflow_addons.custom_ops.image.python import transform as transform_ops
 
-_DTYPES = set([dtypes.uint8, dtypes.int32, dtypes.int64, dtypes.float16,
-               dtypes.float32, dtypes.float64])
+_DTYPES = set([tf.dtypes.uint8, tf.dtypes.int32, tf.dtypes.int64,
+               tf.dtypes.float16, tf.dtypes.float32, tf.dtypes.float64])
 
 
-class ImageOpsTest(test.TestCase):
+class ImageOpsTest(tf.test.TestCase):
     @tf_test_util.run_all_in_graph_and_eager_modes
     def test_compose(self):
         for dtype in _DTYPES:
-            image = constant_op.constant(
+            image = tf.constant(
                 [[1, 1, 1, 0], [1, 0, 0, 0], [1, 1, 1, 0], [0, 0, 0, 0]],
                 dtype=dtype)
             # Rotate counter-clockwise by pi / 2.
             rotation = transform_ops.angles_to_projective_transforms(np.pi / 2,
                                                                      4, 4)
             # Translate right by 1 (the transformation matrix is always inverted,
             # hence the -1).
-            translation = constant_op.constant(
+            translation = tf.constant(
                 [1, 0, -1, 0, 1, 0, 0, 0],
-                dtype=dtypes.float32)
+                dtype=tf.dtypes.float32)
             composed = transform_ops.compose_transforms(rotation, translation)
             image_transformed = transform_ops.transform(image, composed)
             self.assertAllEqual(
@@ -56,31 +53,31 @@ def test_compose(self):
     @tf_test_util.run_all_in_graph_and_eager_modes
     def test_extreme_projective_transform(self):
         for dtype in _DTYPES:
-            image = constant_op.constant(
+            image = tf.constant(
                 [[1, 0, 1, 0], [0, 1, 0, 1], [1, 0, 1, 0], [0, 1, 0, 1]],
                 dtype=dtype)
-            transformation = constant_op.constant(
-                [1, 0, 0, 0, 1, 0, -1, 0], dtypes.float32)
+            transformation = tf.constant(
+                [1, 0, 0, 0, 1, 0, -1, 0], tf.dtypes.float32)
             image_transformed = transform_ops.transform(image, transformation)
             self.assertAllEqual(
                 [[1, 0, 0, 0], [0, 0, 0, 0], [1, 0, 0, 0],
                  [0, 0, 0, 0]], image_transformed)
 
     @tf_test_util.run_all_in_graph_and_eager_modes
     def test_transform_static_output_shape(self):
-        image = constant_op.constant([[1., 2.], [3., 4.]])
+        image = tf.constant([[1., 2.], [3., 4.]])
         result = transform_ops.transform(
             image,
-            random_ops.random_uniform(
+            tf.random.uniform(
                 [8], -1, 1),
-            output_shape=constant_op.constant([3, 5]))
+            output_shape=tf.constant([3, 5]))
         self.assertAllEqual([3, 5], result.shape)
 
     def _test_grad(self, shape_to_test):
         with self.cached_session():
             test_image_shape = shape_to_test
             test_image = np.random.randn(*test_image_shape)
-            test_image_tensor = constant_op.constant(test_image,
+            test_image_tensor = tf.constant(test_image,
                                                      shape=test_image_shape)
             test_transform = transform_ops.angles_to_projective_transforms(
                 np.pi / 2, 4, 4)
@@ -99,7 +96,7 @@ def _test_grad_different_shape(self, input_shape, output_shape):
         with self.cached_session():
             test_image_shape = input_shape
             test_image = np.random.randn(*test_image_shape)
-            test_image_tensor = constant_op.constant(test_image,
+            test_image_tensor = tf.constant(test_image,
                                                      shape=test_image_shape)
             test_transform = transform_ops.angles_to_projective_transforms(
                 np.pi / 2, 4, 4)
@@ -134,19 +131,19 @@ def test_grad(self):
     @tf_test_util.run_all_in_graph_and_eager_modes
     def test_transform_data_types(self):
         for dtype in _DTYPES:
-            image = constant_op.constant([[1, 2], [3, 4]], dtype=dtype)
+            image = tf.constant([[1, 2], [3, 4]], dtype=dtype)
             with self.test_session(use_gpu=True):
                 self.assertAllEqual(
                     np.array([[4, 4], [4, 4]]).astype(dtype.as_numpy_dtype()),
                     transform_ops.transform(image, [1] * 8))
 
     @tf_test_util.run_all_in_graph_and_eager_modes
     def test_transform_eager(self):
-        image = constant_op.constant([[1., 2.], [3., 4.]])
+        image = tf.constant([[1., 2.], [3., 4.]])
         self.assertAllEqual(
             np.array([[4, 4], [4, 4]]),
             transform_ops.transform(image, [1] * 8))
 
 
 if __name__ == "__main__":
-    test.main()
+    tf.test.main()

tensorflow_addons/custom_ops/text/python/skip_gram_ops.py

@@ -18,20 +18,14 @@
 from __future__ import print_function
 
 import csv
+import tensorflow as tf
 
-from tensorflow.python.framework import dtypes
-from tensorflow.python.framework import load_library
-from tensorflow.python.framework import ops
 from tensorflow.python.framework import random_seed
-from tensorflow.python.ops import array_ops
+from tensorflow.python.framework import ops
 from tensorflow.python.ops import lookup_ops
-from tensorflow.python.ops import math_ops
-from tensorflow.python.ops import random_ops
-from tensorflow.python.platform import gfile
 from tensorflow.python.platform import resource_loader
-from tensorflow.python.training import input as input_ops
 
-skip_gram_ops = load_library.load_op_library(
+skip_gram_ops = tf.load_op_library(
     resource_loader.get_path_to_datafile("_skip_gram_ops.so"))
 
 ops.NotDifferentiable("SkipGramGenerateCandidates")
@@ -197,7 +191,7 @@ def skip_gram_sample(input_tensor,
             batch_capacity = (batch_capacity
                               if (batch_capacity is not None and batch_capacity > 0)
                               else 100 * batch_size)
-            return input_ops.batch(
+            return tf.train.batch(
                 [tokens, labels],
                 batch_size,
                 capacity=batch_capacity,
@@ -209,9 +203,9 @@ def skip_gram_sample(input_tensor,
 def skip_gram_sample_with_text_vocab(input_tensor,
                                      vocab_freq_file,
                                      vocab_token_index=0,
-                                     vocab_token_dtype=dtypes.string,
+                                     vocab_token_dtype=tf.dtypes.string,
                                      vocab_freq_index=1,
-                                     vocab_freq_dtype=dtypes.float64,
+                                     vocab_freq_dtype=tf.dtypes.float64,
                                      vocab_delimiter=",",
                                      vocab_min_count=0,
                                      vocab_subsampling=None,
@@ -330,7 +324,7 @@ def skip_gram_sample_with_text_vocab(input_tensor,
     # vocab terms).
     calculated_corpus_size = 0.0
     vocab_size = 0
-    with gfile.GFile(vocab_freq_file, mode="r") as f:
+    with tf.io.gfile.GFile(vocab_freq_file, mode="r") as f:
         reader = csv.reader(f, delimiter=vocab_delimiter)
         for row in reader:
             if vocab_token_index >= len(row) or vocab_freq_index >= len(row):
@@ -404,16 +398,16 @@ def _filter_input(input_tensor, vocab_freq_table, vocab_min_count,
         # Filters out elements in input_tensor that are not found in
         # vocab_freq_table (table returns a default value of -1 specified above when
         # an element is not found).
-        mask = math_ops.not_equal(freq, vocab_freq_table.default_value)
+        mask = tf.math.not_equal(freq, vocab_freq_table.default_value)
 
         # Filters out elements whose vocab frequencies are less than the threshold.
         if vocab_min_count is not None:
-            cast_threshold = math_ops.cast(vocab_min_count, freq.dtype)
-            mask = math_ops.logical_and(mask,
-                                        math_ops.greater_equal(freq, cast_threshold))
+            cast_threshold = tf.cast(vocab_min_count, freq.dtype)
+            mask = tf.math.logical_and(mask,
+                                        tf.math.greater_equal(freq, cast_threshold))
 
-        input_tensor = array_ops.boolean_mask(input_tensor, mask)
-        freq = array_ops.boolean_mask(freq, mask)
+        input_tensor = tf.boolean_mask(input_tensor, mask)
+        freq = tf.boolean_mask(freq, mask)
 
     if not vocab_subsampling:
         return input_tensor
@@ -428,21 +422,21 @@ def _filter_input(input_tensor, vocab_freq_table, vocab_min_count,
     # tokens).
     with ops.name_scope(
             "subsample_vocab", values=[input_tensor, freq, vocab_subsampling]):
-        corpus_size = math_ops.cast(corpus_size, dtypes.float64)
-        freq = math_ops.cast(freq, dtypes.float64)
-        vocab_subsampling = math_ops.cast(vocab_subsampling, dtypes.float64)
+        corpus_size = tf.cast(corpus_size, tf.dtypes.float64)
+        freq = tf.cast(freq, tf.dtypes.float64)
+        vocab_subsampling = tf.cast(vocab_subsampling, tf.dtypes.float64)
 
         # From tensorflow_models/tutorials/embedding/word2vec_kernels.cc, which is
         # suppose to correlate with Eq. 5 in http://arxiv.org/abs/1310.4546.
-        keep_prob = ((math_ops.sqrt(freq /
+        keep_prob = ((tf.math.sqrt(freq /
                                     (vocab_subsampling * corpus_size)) + 1.0) *
                      (vocab_subsampling * corpus_size / freq))
-        random_prob = random_ops.random_uniform(
-            array_ops.shape(freq),
+        random_prob = tf.random.uniform(
+            tf.shape(freq),
             minval=0,
             maxval=1,
-            dtype=dtypes.float64,
+            dtype=tf.dtypes.float64,
             seed=seed)
 
-        mask = math_ops.less_equal(random_prob, keep_prob)
-        return array_ops.boolean_mask(input_tensor, mask)
+        mask = tf.math.less_equal(random_prob, keep_prob)
+        return tf.boolean_mask(input_tensor, mask)