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))