From 6a9a1b0c28f0efa10b468a53b7c2d365d7fdf49c Mon Sep 17 00:00:00 2001 From: "A. Unique TensorFlower" Date: Wed, 19 Aug 2026 09:19:24 -0700 Subject: [PATCH] Fix sample weight generation in MIA tests and add test sharding to single_layer_softmax_test. * Fix MIA sample weight generation: Replace `rng.randn` with `rng.rand` in `membership_inference_attack_test.py` to ensure non-negative sample weights required by `RandomForestClassifier`'s bootstrap sampling logic. * Add sharding to `single_layer_softmax_test`: Migrate to `absltest` and set `shard_count = 4` in the BUILD file to execute parameterized test cases in parallel and reduce test runtime. PiperOrigin-RevId: 967247763 --- .../privacy/logistic_regression/BUILD | 1 + .../single_layer_softmax_test.py | 19 +++++++++++-------- .../membership_inference_attack_test.py | 10 ++++++---- 3 files changed, 18 insertions(+), 12 deletions(-) diff --git a/tensorflow_privacy/privacy/logistic_regression/BUILD b/tensorflow_privacy/privacy/logistic_regression/BUILD index 3161de3d..4feccc3c 100644 --- a/tensorflow_privacy/privacy/logistic_regression/BUILD +++ b/tensorflow_privacy/privacy/logistic_regression/BUILD @@ -56,6 +56,7 @@ py_test( name = "single_layer_softmax_test", size = "medium", srcs = ["single_layer_softmax_test.py"], + shard_count = 4, deps = [ ":datasets", ":single_layer_softmax", diff --git a/tensorflow_privacy/privacy/logistic_regression/single_layer_softmax_test.py b/tensorflow_privacy/privacy/logistic_regression/single_layer_softmax_test.py index cd9db993..07f4e12e 100644 --- a/tensorflow_privacy/privacy/logistic_regression/single_layer_softmax_test.py +++ b/tensorflow_privacy/privacy/logistic_regression/single_layer_softmax_test.py @@ -12,8 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest - +from absl.testing import absltest from absl.testing import parameterized from tensorflow_privacy.privacy.logistic_regression import datasets from tensorflow_privacy.privacy.logistic_regression import single_layer_softmax @@ -27,13 +26,17 @@ class SingleLayerSoftmaxTest(parameterized.TestCase): (10000, 1000, 3, 40, 4, 0.1), (10000, 1000, 4, 40, 4, 0.1), ) - def test_single_layer_softmax(self, num_train, num_test, dimension, epochs, - num_classes, tolerance): - (train_dataset, test_dataset) = datasets.synthetic_linearly_separable_data( - num_train, num_test, dimension, num_classes) + def test_single_layer_softmax( + self, num_train, num_test, dimension, epochs, num_classes, tolerance + ): + train_dataset, test_dataset = datasets.synthetic_linearly_separable_data( + num_train, num_test, dimension, num_classes + ) _, accuracy = single_layer_softmax.single_layer_softmax_classifier( - train_dataset, test_dataset, epochs, num_classes, 'sgd') + train_dataset, test_dataset, epochs, num_classes, 'sgd' + ) self.assertAlmostEqual(accuracy[-1], 1, delta=tolerance) + if __name__ == '__main__': - unittest.main() + absltest.main() diff --git a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/membership_inference_attack_test.py b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/membership_inference_attack_test.py index 4644814d..a0f84317 100644 --- a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/membership_inference_attack_test.py +++ b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/membership_inference_attack_test.py @@ -83,8 +83,9 @@ def get_multilabel_test_input_with_sample_weights(n_train, n_test): logits_test=rng.randn(n_test, num_classes) + 0.2, labels_train=get_multihot_labels_for_test(n_train, num_classes), labels_test=get_multihot_labels_for_test(n_test, num_classes), - sample_weight_train=rng.randn(n_train, 1), - sample_weight_test=rng.randn(n_test, 1)) + sample_weight_train=rng.rand(n_train, 1), + sample_weight_test=rng.rand(n_test, 1), + ) def get_test_input_logits_only(n_train, n_test): @@ -101,8 +102,9 @@ def get_test_input_logits_only_with_sample_weights(n_train, n_test): return AttackInputData( logits_train=rng.randn(n_train, 5) + 0.2, logits_test=rng.randn(n_test, 5) + 0.2, - sample_weight_train=rng.randn(n_train, 1), - sample_weight_test=rng.randn(n_test, 1)) + sample_weight_train=rng.rand(n_train, 1), + sample_weight_test=rng.rand(n_test, 1), + ) class MockTrainedAttacker(object):