Fix MOT TAP test failres.

PiperOrigin-RevId: 341087125

Author: liyunlu0618

tensorflow_model_optimization/python/core/clustering/keras/cluster_integration_test.py

@@ -171,13 +171,13 @@ def testEndToEndDeepLayer(self):
 
     def clusters_check(stripped_model):
       # inner dense layer
-      weights_as_list = stripped_model._layers[1]._layers[1].trainable_weights[0].\
+      weights_as_list = stripped_model.submodules[1].trainable_weights[0].\
         numpy().flatten()
       unique_weights = set(weights_as_list)
       self.assertLessEqual(len(unique_weights), self.params["number_of_clusters"])
 
       # outer dense layer
-      weights_as_list = stripped_model._layers[2].trainable_weights[0].\
+      weights_as_list = stripped_model.submodules[4].trainable_weights[0].\
         numpy().flatten()
       unique_weights = set(weights_as_list)
       self.assertLessEqual(len(unique_weights), self.params["number_of_clusters"])
@@ -199,20 +199,20 @@ def testEndToEndDeepLayer2(self):
 
     def clusters_check(stripped_model):
       # first inner dense layer
-      weights_as_list = stripped_model._layers[1]._layers[1].trainable_weights[0].\
+      weights_as_list = stripped_model.submodules[1].trainable_weights[0].\
         numpy().flatten()
       unique_weights = set(weights_as_list)
       self.assertLessEqual(len(unique_weights), self.params["number_of_clusters"])
 
       # second inner dense layer
-      weights_as_list = stripped_model._layers[1]._layers[1]._layers[1].\
+      weights_as_list = stripped_model.submodules[4].\
         trainable_weights[0].\
         numpy().flatten()
       unique_weights = set(weights_as_list)
       self.assertLessEqual(len(unique_weights), self.params["number_of_clusters"])
 
       # outer dense layer
-      weights_as_list = stripped_model._layers[2].trainable_weights[0].\
+      weights_as_list = stripped_model.submodules[7].trainable_weights[0].\
         numpy().flatten()
       unique_weights = set(weights_as_list)
       self.assertLessEqual(len(unique_weights), self.params["number_of_clusters"])

tensorflow_model_optimization/python/core/sparsity/keras/prune_test.py

@@ -116,7 +116,7 @@ def _validate_pruned_layer(self, original_layer, wrapped_layer):
   @staticmethod
   def _count_pruned_layers(model):
     count = 0
-    for layer in model._layers:
+    for layer in model.submodules:
       if isinstance(layer, pruning_wrapper.PruneLowMagnitude):
         count += 1
     return count