The previous version uses tf.nest.map_structure to apply add_noise to a tf.RaggedTensor. This causes a bug when used in tensorflow federated because tf.nest.map_structure will also map add_noise to the tensor for shape information in tf.RaggedTensor. This causes failure when tff conducts automatic type conversion.

Also use fixed random seed to avoid flaky timeouts and testing failures.

PiperOrigin-RevId: 384573740

Author: tensorflower-gardener

tensorflow_privacy/privacy/dp_query/tree_aggregation_query.py

@@ -25,10 +25,10 @@
 """
 import distutils
 import math
-import attr
+from typing import Optional
 
+import attr
 import tensorflow as tf
-
 from tensorflow_privacy.privacy.dp_query import dp_query
 from tensorflow_privacy.privacy.dp_query import tree_aggregation
 
@@ -442,16 +442,20 @@ def _loop_body(i, h):
   return tree
 
 
-def _get_add_noise(stddev):
+def _get_add_noise(stddev, seed: int = None):
   """Utility function to decide which `add_noise` to use according to tf version."""
   if distutils.version.LooseVersion(
       tf.__version__) < distutils.version.LooseVersion('2.0.0'):
 
+    # The seed should be only used for testing purpose.
+    if seed is not None:
+      tf.random.set_seed(seed)
+
     def add_noise(v):
       return v + tf.random.normal(
           tf.shape(input=v), stddev=stddev, dtype=v.dtype)
   else:
-    random_normal = tf.random_normal_initializer(stddev=stddev)
+    random_normal = tf.random_normal_initializer(stddev=stddev, seed=seed)
 
     def add_noise(v):
       return v + tf.cast(random_normal(tf.shape(input=v)), dtype=v.dtype)
@@ -478,17 +482,16 @@ class GlobalState(object):
     """Class defining global state for `CentralTreeSumQuery`.
 
     Attributes:
-      stddev: The stddev of the noise added to each node in the tree.
-      arity: The branching factor of the tree (i.e. the number of children each
-        internal node has).
       l1_bound: An upper bound on the L1 norm of the input record. This is
         needed to bound the sensitivity and deploy differential privacy.
     """
-    stddev = attr.ib()
-    arity = attr.ib()
     l1_bound = attr.ib()
 
-  def __init__(self, stddev: float, arity: int = 2, l1_bound: int = 10):
+  def __init__(self,
+               stddev: float,
+               arity: int = 2,
+               l1_bound: int = 10,
+               seed: Optional[int] = None):
     """Initializes the `CentralTreeSumQuery`.
 
     Args:
@@ -497,15 +500,17 @@ def __init__(self, stddev: float, arity: int = 2, l1_bound: int = 10):
       arity: The branching factor of the tree.
       l1_bound: An upper bound on the L1 norm of the input record. This is
         needed to bound the sensitivity and deploy differential privacy.
+      seed: Random seed to generate Gaussian noise. Defaults to `None`. Only for
+        test purpose.
     """
     self._stddev = stddev
     self._arity = arity
     self._l1_bound = l1_bound
+    self._seed = seed
 
   def initial_global_state(self):
     """Implements `tensorflow_privacy.DPQuery.initial_global_state`."""
-    return CentralTreeSumQuery.GlobalState(
-        stddev=self._stddev, arity=self._arity, l1_bound=self._l1_bound)
+    return CentralTreeSumQuery.GlobalState(l1_bound=self._l1_bound)
 
   def derive_sample_params(self, global_state):
     """Implements `tensorflow_privacy.DPQuery.derive_sample_params`."""
@@ -536,10 +541,9 @@ def get_noised_result(self, sample_state, global_state):
       The jth node on the ith layer of the tree can be accessed by tree[i][j]
       where tree is the returned value.
     """
-    add_noise = _get_add_noise(self._stddev)
-    tree = _build_tree_from_leaf(sample_state, global_state.arity)
-    return tf.nest.map_structure(
-        add_noise, tree, expand_composites=True), global_state
+    add_noise = _get_add_noise(self._stddev, self._seed)
+    tree = _build_tree_from_leaf(sample_state, self._arity)
+    return tf.map_fn(add_noise, tree), global_state
 
 
 class DistributedTreeSumQuery(dp_query.SumAggregationDPQuery):
@@ -577,18 +581,25 @@ class GlobalState(object):
     arity = attr.ib()
     l1_bound = attr.ib()
 
-  def __init__(self, stddev: float, arity: int = 2, l1_bound: int = 10):
+  def __init__(self,
+               stddev: float,
+               arity: int = 2,
+               l1_bound: int = 10,
+               seed: Optional[int] = None):
     """Initializes the `DistributedTreeSumQuery`.
 
     Args:
       stddev: The stddev of the noise added to each node in the tree.
       arity: The branching factor of the tree.
       l1_bound: An upper bound on the L1 norm of the input record. This is
         needed to bound the sensitivity and deploy differential privacy.
+      seed: Random seed to generate Gaussian noise. Defaults to `None`. Only for
+        test purpose.
     """
     self._stddev = stddev
     self._arity = arity
     self._l1_bound = l1_bound
+    self._seed = seed
 
   def initial_global_state(self):
     """Implements `tensorflow_privacy.DPQuery.initial_global_state`."""
@@ -628,9 +639,9 @@ def preprocess_record(self, params, record):
                                                     use_norm=l1_norm)
     preprocessed_record = preprocessed_record[0]
 
-    add_noise = _get_add_noise(self._stddev)
+    add_noise = _get_add_noise(self._stddev, self._seed)
     tree = _build_tree_from_leaf(preprocessed_record, arity)
-    noisy_tree = tf.nest.map_structure(add_noise, tree, expand_composites=True)
+    noisy_tree = tf.map_fn(add_noise, tree)
 
     # The following codes reshape the output vector so the output shape of can
     # be statically inferred. This is useful when used with

tensorflow_privacy/privacy/dp_query/tree_aggregation_query_test.py

@@ -502,21 +502,15 @@ def test_get_noised_result(self, arity, record, expected_tree):
       ('stddev_0_1', 0.1, tf.constant([1, 0], dtype=tf.int32), [1., 1., 0.]),
   )
   def test_get_noised_result_with_noise(self, stddev, record, expected_tree):
-    query = tree_aggregation_query.CentralTreeSumQuery(stddev=stddev)
+    query = tree_aggregation_query.CentralTreeSumQuery(stddev=stddev, seed=0)
     global_state = query.initial_global_state()
     params = query.derive_sample_params(global_state)
     preprocessed_record = query.preprocess_record(params, record)
-    sample_state_list = []
-    for _ in range(1000):
-      sample_state, _ = query.get_noised_result(preprocessed_record,
-                                                global_state)
-      sample_state_list.append(sample_state.flat_values.numpy())
-    expectation = np.mean(sample_state_list, axis=0)
-    variance = np.std(sample_state_list, axis=0)
-
-    self.assertAllClose(expectation, expected_tree, rtol=3 * stddev, atol=1e-4)
+
+    sample_state, _ = query.get_noised_result(preprocessed_record, global_state)
+
     self.assertAllClose(
-        variance, np.ones(len(variance)) * stddev, rtol=0.1, atol=1e-4)
+        sample_state.flat_values, expected_tree, atol=3 * stddev)
 
   @parameterized.named_parameters(
       ('binary_test_int', 2, tf.constant([10, 10, 0, 0], dtype=tf.int32),
@@ -556,8 +550,7 @@ def test_initial_global_state_type(self):
   def test_derive_sample_params(self):
     query = tree_aggregation_query.DistributedTreeSumQuery(stddev=NOISE_STD)
     global_state = query.initial_global_state()
-    stddev, arity, l1_bound = query.derive_sample_params(
-        global_state)
+    stddev, arity, l1_bound = query.derive_sample_params(global_state)
     self.assertAllClose(stddev, NOISE_STD)
     self.assertAllClose(arity, 2)
     self.assertAllClose(l1_bound, 10)
@@ -587,21 +580,14 @@ def test_preprocess_record(self, arity, record, expected_tree):
       ('stddev_0_1', 0.1, tf.constant([1, 0], dtype=tf.int32), [1., 1., 0.]),
   )
   def test_preprocess_record_with_noise(self, stddev, record, expected_tree):
-    query = tree_aggregation_query.DistributedTreeSumQuery(stddev=stddev)
+    query = tree_aggregation_query.DistributedTreeSumQuery(
+        stddev=stddev, seed=0)
     global_state = query.initial_global_state()
     params = query.derive_sample_params(global_state)
 
-    preprocessed_record_list = []
-    for _ in range(1000):
-      preprocessed_record = query.preprocess_record(params, record)
-      preprocessed_record_list.append(preprocessed_record.numpy())
-
-    expectation = np.mean(preprocessed_record_list, axis=0)
-    variance = np.std(preprocessed_record_list, axis=0)
+    preprocessed_record = query.preprocess_record(params, record)
 
-    self.assertAllClose(expectation, expected_tree, rtol=3 * stddev, atol=1e-4)
-    self.assertAllClose(
-        variance, np.ones(len(variance)) * stddev, rtol=0.1, atol=1e-4)
+    self.assertAllClose(preprocessed_record, expected_tree, atol=3 * stddev)
 
   @parameterized.named_parameters(
       ('binary_test_int', 2, tf.constant([10, 10, 0, 0], dtype=tf.int32),