Fix issue with importing tensorflow.compat.v1.

PiperOrigin-RevId: 300175680

tensorflow_privacy/privacy/analysis/privacy_ledger.py

@@ -111,7 +111,7 @@ def record_sum_query(self, l2_norm_bound, noise_stddev):
 
     def _do_record_query():
       with tf.control_dependencies(
-          [tf.compat.v1.assign(self._query_count, self._query_count + 1)]):
+          [tf.assign(self._query_count, self._query_count + 1)]):
         return self._query_buffer.append(
             [self._sample_count, l2_norm_bound, noise_stddev])
 
@@ -120,14 +120,14 @@ def _do_record_query():
   def finalize_sample(self):
     """Finalizes sample and records sample ledger entry."""
     with tf.control_dependencies([
-        tf.compat.v1.assign(self._sample_var, [
+        tf.assign(self._sample_var, [
             self._population_size, self._selection_probability,
             self._query_count
         ])
     ]):
       with tf.control_dependencies([
-          tf.compat.v1.assign(self._sample_count, self._sample_count + 1),
-          tf.compat.v1.assign(self._query_count, 0)
+          tf.assign(self._sample_count, self._sample_count + 1),
+          tf.assign(self._query_count, 0)
       ]):
         return self._sample_buffer.append(self._sample_var)
 

tensorflow_privacy/privacy/analysis/privacy_ledger_test.py

@@ -25,7 +25,7 @@
 from tensorflow_privacy.privacy.dp_query import nested_query
 from tensorflow_privacy.privacy.dp_query import test_utils
 
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 class PrivacyLedgerTest(tf.test.TestCase):
@@ -63,8 +63,8 @@ def test_sum_query(self):
         query, population_size, selection_probability)
 
     # First sample.
-    tf.compat.v1.assign(population_size, 10)
-    tf.compat.v1.assign(selection_probability, 0.1)
+    tf.assign(population_size, 10)
+    tf.assign(selection_probability, 0.1)
     test_utils.run_query(query, [record1, record2])
 
     expected_queries = [[10.0, 0.0]]
@@ -75,8 +75,8 @@ def test_sum_query(self):
     self.assertAllClose(sample_1.queries, expected_queries)
 
     # Second sample.
-    tf.compat.v1.assign(population_size, 20)
-    tf.compat.v1.assign(selection_probability, 0.2)
+    tf.assign(population_size, 20)
+    tf.assign(selection_probability, 0.2)
     test_utils.run_query(query, [record1, record2])
 
     formatted = query.ledger.get_formatted_ledger_eager()
@@ -106,8 +106,8 @@ def test_nested_query(self):
     record2 = [5.0, [1.0, 2.0]]
 
     # First sample.
-    tf.compat.v1.assign(population_size, 10)
-    tf.compat.v1.assign(selection_probability, 0.1)
+    tf.assign(population_size, 10)
+    tf.assign(selection_probability, 0.1)
     test_utils.run_query(query, [record1, record2])
 
     expected_queries = [[4.0, 2.0], [5.0, 1.0]]
@@ -118,8 +118,8 @@ def test_nested_query(self):
     self.assertAllClose(sorted(sample_1.queries), sorted(expected_queries))
 
     # Second sample.
-    tf.compat.v1.assign(population_size, 20)
-    tf.compat.v1.assign(selection_probability, 0.2)
+    tf.assign(population_size, 20)
+    tf.assign(selection_probability, 0.2)
     test_utils.run_query(query, [record1, record2])
 
     formatted = query.ledger.get_formatted_ledger_eager()

tensorflow_privacy/privacy/analysis/tensor_buffer.py

@@ -50,10 +50,10 @@ def __init__(self, capacity, shape, dtype=tf.int32, name=None):
       raise ValueError('Shape cannot be scalar.')
     shape = [capacity] + shape
 
-    with tf.compat.v1.variable_scope(self._name):
+    with tf.variable_scope(self._name):
       # We need to use a placeholder as the initial value to allow resizing.
-      self._buffer = tf.compat.v1.Variable(
-          initial_value=tf.compat.v1.placeholder_with_default(
+      self._buffer = tf.Variable(
+          initial_value=tf.placeholder_with_default(
               tf.zeros(shape, dtype), shape=None),
           trainable=False,
           name='buffer',
@@ -82,18 +82,18 @@ def _double_capacity():
       padding = tf.zeros_like(self._buffer, self._buffer.dtype)
       new_buffer = tf.concat([self._buffer, padding], axis=0)
       if tf.executing_eagerly():
-        with tf.compat.v1.variable_scope(self._name, reuse=True):
-          self._buffer = tf.compat.v1.get_variable(
+        with tf.variable_scope(self._name, reuse=True):
+          self._buffer = tf.get_variable(
               name='buffer',
               dtype=self._dtype,
               initializer=new_buffer,
               trainable=False)
-          return self._buffer, tf.compat.v1.assign(
+          return self._buffer, tf.assign(
               self._capacity, tf.multiply(self._capacity, 2))
       else:
-        return tf.compat.v1.assign(
+        return tf.assign(
             self._buffer, new_buffer,
-            validate_shape=False), tf.compat.v1.assign(
+            validate_shape=False), tf.assign(
                 self._capacity, tf.multiply(self._capacity, 2))
 
     update_buffer, update_capacity = tf.cond(
@@ -103,18 +103,18 @@ def _double_capacity():
 
     with tf.control_dependencies([update_buffer, update_capacity]):
       with tf.control_dependencies([
-          tf.compat.v1.assert_less(
+          tf.assert_less(
               self._current_size,
               self._capacity,
               message='Appending past end of TensorBuffer.'),
-          tf.compat.v1.assert_equal(
+          tf.assert_equal(
               tf.shape(input=value),
               tf.shape(input=self._buffer)[1:],
               message='Appending value of inconsistent shape.')
       ]):
         with tf.control_dependencies(
-            [tf.compat.v1.assign(self._buffer[self._current_size, :], value)]):
-          return tf.compat.v1.assign_add(self._current_size, 1)
+            [tf.assign(self._buffer[self._current_size, :], value)]):
+          return tf.assign_add(self._current_size, 1)
 
   @property
   def values(self):

tensorflow_privacy/privacy/analysis/tensor_buffer_test_eager.py

@@ -21,7 +21,7 @@
 
 from tensorflow_privacy.privacy.analysis import tensor_buffer
 
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 class TensorBufferTest(tf.test.TestCase):

tensorflow_privacy/privacy/analysis/tensor_buffer_test_graph.py

@@ -38,7 +38,7 @@ def test_noresize(self):
           values = my_buffer.values
           current_size = my_buffer.current_size
           capacity = my_buffer.capacity
-      self.evaluate(tf.compat.v1.global_variables_initializer())
+      self.evaluate(tf.global_variables_initializer())
 
       v, cs, cap = sess.run([values, current_size, capacity])
       self.assertAllEqual(v, [value1, value2])
@@ -60,7 +60,7 @@ def test_resize(self):
             values = my_buffer.values
             current_size = my_buffer.current_size
             capacity = my_buffer.capacity
-      self.evaluate(tf.compat.v1.global_variables_initializer())
+      self.evaluate(tf.global_variables_initializer())
 
       v, cs, cap = sess.run([values, current_size, capacity])
       self.assertAllEqual(v, [value1, value2, value3])

tensorflow_privacy/privacy/dp_query/gaussian_query.py

@@ -96,7 +96,7 @@ def add_noise(v):
         return v + tf.random.normal(
             tf.shape(input=v), stddev=global_state.stddev)
     else:
-      random_normal = tf.compat.v1.random_normal_initializer(
+      random_normal = tf.random_normal_initializer(
           stddev=global_state.stddev)
 
       def add_noise(v):

tensorflow_privacy/privacy/dp_query/gaussian_query_test.py

@@ -59,14 +59,13 @@ def test_gaussian_sum_with_changing_clip_no_noise(self):
       record2 = tf.constant([4.0, -3.0])  # Not clipped.
 
       l2_norm_clip = tf.Variable(5.0)
-      l2_norm_clip_placeholder = tf.compat.v1.placeholder(tf.float32)
-      assign_l2_norm_clip = tf.compat.v1.assign(l2_norm_clip,
-                                                l2_norm_clip_placeholder)
+      l2_norm_clip_placeholder = tf.placeholder(tf.float32)
+      assign_l2_norm_clip = tf.assign(l2_norm_clip, l2_norm_clip_placeholder)
       query = gaussian_query.GaussianSumQuery(
           l2_norm_clip=l2_norm_clip, stddev=0.0)
       query_result, _ = test_utils.run_query(query, [record1, record2])
 
-      self.evaluate(tf.compat.v1.global_variables_initializer())
+      self.evaluate(tf.global_variables_initializer())
       result = sess.run(query_result)
       expected = [1.0, 1.0]
       self.assertAllClose(result, expected)

tensorflow_privacy/privacy/dp_query/quantile_adaptive_clip_sum_query_test.py

@@ -27,7 +27,7 @@
 from tensorflow_privacy.privacy.dp_query import quantile_adaptive_clip_sum_query
 from tensorflow_privacy.privacy.dp_query import test_utils
 
-tf.compat.v1.enable_eager_execution()
+tf.enable_eager_execution()
 
 
 class QuantileAdaptiveClipSumQueryTest(
@@ -323,7 +323,7 @@ def test_adaptation_all_equal(self, start_low, geometric):
     global_state = query.initial_global_state()
 
     for t in range(50):
-      tf.compat.v1.assign(learning_rate, 1.0 / np.sqrt(t + 1))
+      tf.assign(learning_rate, 1.0 / np.sqrt(t + 1))
       _, global_state = test_utils.run_query(query, records, global_state)
 
       actual_clip = global_state.sum_state.l2_norm_clip
@@ -350,8 +350,8 @@ def test_ledger(self):
         query, population_size, selection_probability)
 
     # First sample.
-    tf.compat.v1.assign(population_size, 10)
-    tf.compat.v1.assign(selection_probability, 0.1)
+    tf.assign(population_size, 10)
+    tf.assign(selection_probability, 0.1)
     _, global_state = test_utils.run_query(query, [record1, record2])
 
     expected_queries = [[10.0, 10.0], [0.5, 0.0]]
@@ -362,8 +362,8 @@ def test_ledger(self):
     self.assertAllClose(sample_1.queries, expected_queries)
 
     # Second sample.
-    tf.compat.v1.assign(population_size, 20)
-    tf.compat.v1.assign(selection_probability, 0.2)
+    tf.assign(population_size, 20)
+    tf.assign(selection_probability, 0.2)
     test_utils.run_query(query, [record1, record2], global_state)
 
     formatted = query.ledger.get_formatted_ledger_eager()

tensorflow_privacy/privacy/optimizers/dp_optimizer.py

@@ -27,9 +27,9 @@
 
 def make_optimizer_class(cls):
   """Constructs a DP optimizer class from an existing one."""
-  parent_code = tf.compat.v1.train.Optimizer.compute_gradients.__code__
+  parent_code = tf.train.Optimizer.compute_gradients.__code__
   child_code = cls.compute_gradients.__code__
-  GATE_OP = tf.compat.v1.train.Optimizer.GATE_OP  # pylint: disable=invalid-name
+  GATE_OP = tf.train.Optimizer.GATE_OP  # pylint: disable=invalid-name
   if child_code is not parent_code:
     logging.warning(
         'WARNING: Calling make_optimizer_class() on class %s that overrides '
@@ -146,8 +146,8 @@ def process_microbatch(i, sample_state):
 
         if var_list is None:
           var_list = (
-              tf.compat.v1.trainable_variables() + tf.compat.v1.get_collection(
-                  tf.compat.v1.GraphKeys.TRAINABLE_RESOURCE_VARIABLES))
+              tf.trainable_variables() + tf.get_collection(
+                  tf.GraphKeys.TRAINABLE_RESOURCE_VARIABLES))
 
         sample_state = self._dp_sum_query.initial_sample_state(var_list)
 
@@ -213,9 +213,9 @@ def ledger(self):
 
   return DPGaussianOptimizerClass
 
-AdagradOptimizer = tf.compat.v1.train.AdagradOptimizer
-AdamOptimizer = tf.compat.v1.train.AdamOptimizer
-GradientDescentOptimizer = tf.compat.v1.train.GradientDescentOptimizer
+AdagradOptimizer = tf.train.AdagradOptimizer
+AdamOptimizer = tf.train.AdamOptimizer
+GradientDescentOptimizer = tf.train.GradientDescentOptimizer
 
 DPAdagradOptimizer = make_optimizer_class(AdagradOptimizer)
 DPAdamOptimizer = make_optimizer_class(AdamOptimizer)

tensorflow_privacy/privacy/optimizers/dp_optimizer_eager_test.py

@@ -30,7 +30,7 @@
 class DPOptimizerEagerTest(tf.test.TestCase, parameterized.TestCase):
 
   def setUp(self):
-    tf.compat.v1.enable_eager_execution()
+    tf.enable_eager_execution()
     super(DPOptimizerEagerTest, self).setUp()
 
   def _loss_fn(self, val0, val1):
@@ -64,7 +64,7 @@ def testBaseline(self, cls, num_microbatches, expected_answer):
           num_microbatches=num_microbatches,
           learning_rate=2.0)
 
-      self.evaluate(tf.compat.v1.global_variables_initializer())
+      self.evaluate(tf.global_variables_initializer())
       # Fetch params to validate initial values
       self.assertAllClose([1.0, 2.0], self.evaluate(var0))
 
@@ -89,7 +89,7 @@ def testClippingNorm(self, cls):
 
       opt = cls(dp_sum_query, num_microbatches=1, learning_rate=2.0)
 
-      self.evaluate(tf.compat.v1.global_variables_initializer())
+      self.evaluate(tf.global_variables_initializer())
       # Fetch params to validate initial values
       self.assertAllClose([0.0, 0.0], self.evaluate(var0))
 
@@ -113,7 +113,7 @@ def testNoiseMultiplier(self, cls):
 
       opt = cls(dp_sum_query, num_microbatches=1, learning_rate=2.0)
 
-      self.evaluate(tf.compat.v1.global_variables_initializer())
+      self.evaluate(tf.global_variables_initializer())
       # Fetch params to validate initial values
       self.assertAllClose([0.0], self.evaluate(var0))
 

tensorflow_privacy/privacy/optimizers/dp_optimizer_test.py

@@ -63,7 +63,7 @@ def testBaseline(self, cls, num_microbatches, expected_answer):
           num_microbatches=num_microbatches,
           learning_rate=2.0)
 
-      self.evaluate(tf.compat.v1.global_variables_initializer())
+      self.evaluate(tf.global_variables_initializer())
       # Fetch params to validate initial values
       self.assertAllClose([1.0, 2.0], self.evaluate(var0))
 
@@ -87,7 +87,7 @@ def testClippingNorm(self, cls):
 
       opt = cls(dp_sum_query, num_microbatches=1, learning_rate=2.0)
 
-      self.evaluate(tf.compat.v1.global_variables_initializer())
+      self.evaluate(tf.global_variables_initializer())
       # Fetch params to validate initial values
       self.assertAllClose([0.0, 0.0], self.evaluate(var0))
 
@@ -110,7 +110,7 @@ def testNoiseMultiplier(self, cls):
 
       opt = cls(dp_sum_query, num_microbatches=1, learning_rate=2.0)
 
-      self.evaluate(tf.compat.v1.global_variables_initializer())
+      self.evaluate(tf.global_variables_initializer())
       # Fetch params to validate initial values
       self.assertAllClose([0.0], self.evaluate(var0))
 
@@ -126,7 +126,7 @@ def testNoiseMultiplier(self, cls):
   @mock.patch('absl.logging.warning')
   def testComputeGradientsOverrideWarning(self, mock_logging):
 
-    class SimpleOptimizer(tf.compat.v1.train.Optimizer):
+    class SimpleOptimizer(tf.train.Optimizer):
 
       def compute_gradients(self):
         return 0
@@ -153,7 +153,7 @@ def linear_model_fn(features, labels, mode):
           dp_sum_query,
           num_microbatches=1,
           learning_rate=1.0)
-      global_step = tf.compat.v1.train.get_global_step()
+      global_step = tf.train.get_global_step()
       train_op = optimizer.minimize(loss=vector_loss, global_step=global_step)
       return tf.estimator.EstimatorSpec(
           mode=mode, loss=scalar_loss, train_op=train_op)
@@ -167,7 +167,7 @@ def linear_model_fn(features, labels, mode):
                              true_weights) + true_bias + np.random.normal(
                                  scale=0.1, size=(200, 1)).astype(np.float32)
 
-    train_input_fn = tf.compat.v1.estimator.inputs.numpy_input_fn(
+    train_input_fn = tf.estimator.inputs.numpy_input_fn(
         x={'x': train_data},
         y=train_labels,
         batch_size=20,
@@ -200,7 +200,7 @@ def testUnrollMicrobatches(self, cls):
           learning_rate=2.0,
           unroll_microbatches=True)
 
-      self.evaluate(tf.compat.v1.global_variables_initializer())
+      self.evaluate(tf.global_variables_initializer())
       # Fetch params to validate initial values
       self.assertAllClose([1.0, 2.0], self.evaluate(var0))
 
@@ -225,7 +225,7 @@ def testDPGaussianOptimizerClass(self, cls):
           num_microbatches=1,
           learning_rate=2.0)
 
-      self.evaluate(tf.compat.v1.global_variables_initializer())
+      self.evaluate(tf.global_variables_initializer())
       # Fetch params to validate initial values
       self.assertAllClose([0.0], self.evaluate(var0))