From c5f35b3ca1226edc8d82779da668f8cdc8388aa8 Mon Sep 17 00:00:00 2001 From: Zheng Xu Date: Wed, 8 Sep 2021 13:58:53 -0700 Subject: [PATCH] Try to fix flakiness by reducing the number of query construction. PiperOrigin-RevId: 395552656 --- .../privacy/dp_query/tree_aggregation_query_test.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tensorflow_privacy/privacy/dp_query/tree_aggregation_query_test.py b/tensorflow_privacy/privacy/dp_query/tree_aggregation_query_test.py index ace6484..1115f40 100644 --- a/tensorflow_privacy/privacy/dp_query/tree_aggregation_query_test.py +++ b/tensorflow_privacy/privacy/dp_query/tree_aggregation_query_test.py @@ -221,9 +221,9 @@ class TreeCumulativeSumQueryTest(tf.test.TestCase, parameterized.TestCase): _get_noise_generator), ) def test_noisy_cumsum_and_state_update(self, records, value_generator): - num_trials = 200 - record_specs = tf.TensorSpec([]) - records = [tf.constant(r) for r in records] + num_trials, vector_size = 10, 100 + record_specs = tf.TensorSpec([vector_size]) + records = [tf.constant(r, shape=[vector_size]) for r in records] noised_sums = [] for i in range(num_trials): query = tree_aggregation_query.TreeCumulativeSumQuery(