From b607fa381adf3deb8a03020a56dab87cd787e247 Mon Sep 17 00:00:00 2001 From: Ricardo Torres Date: Mon, 1 Jun 2026 07:51:46 -0700 Subject: [PATCH] The `fully_commutes_with_sum` property was incorrectly using `sum()` on the flattened commuting structure, which does not yield a boolean. Changed to use `all()` to ensure the property returns True only if all sub-encoders fully commute with sum. Added type assertions in tests to verify the return type is boolean. PiperOrigin-RevId: 924707489 --- .../internal/tensor_encoding/core/gather_encoder.py | 2 +- .../tensor_encoding/core/gather_encoder_test.py | 12 +++++++++--- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/tensorflow_model_optimization/python/core/internal/tensor_encoding/core/gather_encoder.py b/tensorflow_model_optimization/python/core/internal/tensor_encoding/core/gather_encoder.py index 38b8b2d81..1ce87307c 100644 --- a/tensorflow_model_optimization/python/core/internal/tensor_encoding/core/gather_encoder.py +++ b/tensorflow_model_optimization/python/core/internal/tensor_encoding/core/gather_encoder.py @@ -370,7 +370,7 @@ def input_tensorspec(self): @property def fully_commutes_with_sum(self): # If any element is not True, the whole thing does not fully commute. - return sum(tf.nest.flatten(self._commuting_structure)) + return all(tf.nest.flatten(self._commuting_structure)) @property def state_update_aggregation_modes(self): diff --git a/tensorflow_model_optimization/python/core/internal/tensor_encoding/core/gather_encoder_test.py b/tensorflow_model_optimization/python/core/internal/tensor_encoding/core/gather_encoder_test.py index 26fbed255..be0da871e 100644 --- a/tensorflow_model_optimization/python/core/internal/tensor_encoding/core/gather_encoder_test.py +++ b/tensorflow_model_optimization/python/core/internal/tensor_encoding/core/gather_encoder_test.py @@ -197,13 +197,13 @@ def test_full_commutativity_with_sum(self): encoder = gather_encoder.GatherEncoder.from_encoder( core_encoder.EncoderComposer(test_utils.TimesTwoEncodingStage()).make(), spec) - self.assertTrue(encoder.fully_commutes_with_sum) + self.assertIs(encoder.fully_commutes_with_sum, True) encoder = gather_encoder.GatherEncoder.from_encoder( core_encoder.EncoderComposer( test_utils.TimesTwoEncodingStage()).add_parent( test_utils.TimesTwoEncodingStage(), T2_VALS).make(), spec) - self.assertTrue(encoder.fully_commutes_with_sum) + self.assertIs(encoder.fully_commutes_with_sum, True) encoder = core_encoder.EncoderComposer( test_utils.SignIntFloatEncodingStage()) @@ -212,7 +212,13 @@ def test_full_commutativity_with_sum(self): encoder.add_child(test_utils.TimesTwoEncodingStage(), SIF_FLOATS).add_child( test_utils.PlusOneOverNEncodingStage(), T2_VALS) encoder = gather_encoder.GatherEncoder.from_encoder(encoder.make(), spec) - self.assertFalse(encoder.fully_commutes_with_sum) + self.assertIs(encoder.fully_commutes_with_sum, False) + + encoder = core_encoder.EncoderComposer( + test_utils.TimesTwoEncodingStage()) + encoder.add_child(test_utils.PlusOneEncodingStage(), T2_VALS) + encoder = gather_encoder.GatherEncoder.from_encoder(encoder.make(), spec) + self.assertIs(encoder.fully_commutes_with_sum, False) @tf_test_util.run_all_in_graph_and_eager_modes def test_state_aggregation_modes(self):