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 38b8b2d8..1ce87307 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 26fbed25..be0da871 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):