Add nan / inf tests to the federated_sample aggregation. This aggregation
should preserve the percentage of both in the final sample.
@@ -205,8 +205,27 @@ class FederatedSampleTest(tf.test.TestCase):
return federated_aggregations.federated_sample(value)
value = call_federated_sample([1.0] * 100 + [0.0] * 100)
+ self.assertLen(value, 100)
self.assertAlmostEqual(len(np.nonzero(value)[0]), 50, delta=10)
+ def test_federated_sample_preserves_nan_percentage(self):
+
+ @tff.federated_computation(tff.FederatedType(tf.float32, tff.CLIENTS))
+ def call_federated_sample(value):
+ return federated_aggregations.federated_sample(value)
+
+ value = call_federated_sample([1.0] * 100 + [np.nan] * 100)
+ self.assertAlmostEqual(np.count_nonzero(np.isnan(value)), 50, delta=10)
+
+ def test_federated_sample_preserves_inf_percentage(self):
+
+ @tff.federated_computation(tff.FederatedType(tf.float32, tff.CLIENTS))
+ def call_federated_sample(value):
+ return federated_aggregations.federated_sample(value)
+
+ value = call_federated_sample([1.0] * 100 + [np.inf] * 100)
+ self.assertAlmostEqual(np.count_nonzero(np.isinf(value)), 50, delta=10)
+
if __name__ == '__main__':
absltest.main()
