def test_build_dataset_reduce_fn(self, simulation, reduce_fn):
dataset_reduce_fn = dp_fedavg._build_dataset_reduce_fn(simulation)
self.assertIs(dataset_reduce_fn, reduce_fn)
ds = tf.data.Dataset.range(10, output_type=tf.int32)
total_sum = dataset_reduce_fn(
reduce_fn=lambda x, y: x + y, dataset=ds, initial_state_fn=lambda: 0)
self.assertEqual(total_sum, np.int32(45))
def test_build_dataset_reduce_fn_tuple(self, simulation, reduce_fn):
dataset_reduce_fn = dp_fedavg._build_dataset_reduce_fn(simulation)
self.assertIs(dataset_reduce_fn, reduce_fn)
ds = tf.data.Dataset.range(
10, output_type=tf.float32).map(lambda x: 0.1 * x)
total_cnt, total_sum = dataset_reduce_fn(
reduce_fn=lambda x, y: (x[0] + 1, x[1] + y),
dataset=ds,
initial_state_fn=lambda: (tf.constant(0), tf.constant(0.1)))
self.assertEqual(total_cnt, np.float32(10))
self.assertEqual(total_sum, np.float32(4.6))