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

Author: ratgr

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):

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):