Add basic TF ops to the supported list of PruneForLatencyOnXNNPack policy

PiperOrigin-RevId: 376909045

Author: tensorflower-gardener

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

@@ -217,6 +217,9 @@ def _check_layer_support(self, layer):
       return activations.serialize(layer.activation) in ('relu', 'relu6',
                                                          'leaky_relu', 'elu',
                                                          'sigmoid')
+    elif layer.__class__.__name__ == 'TFOpLambda':
+      return layer.function in (tf.identity, tf.__operators__.add, tf.math.add,
+                                tf.math.subtract, tf.math.multiply)
     return False
 
   def ensure_model_supports_pruning(self, model):

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

@@ -301,6 +301,30 @@ def testPruneFunctionalModelAfterCloneForLatencyOnXNNPackPolicy(self):
     )
     self.assertEqual(self._count_pruned_layers(pruned_model), 1)
 
+  def testFunctionalModelWithTFOpsForLatencyOnXNNPackPolicy(self):
+    i = keras.Input(shape=(8, 8, 3))
+    x = layers.ZeroPadding2D(padding=1)(i)
+    x = layers.Conv2D(
+        filters=16,
+        kernel_size=(3, 3),
+        strides=(2, 2),
+        padding='valid',
+    )(x)
+    residual = layers.Conv2D(filters=16, kernel_size=[1, 1])(x)
+    x = x + residual
+    x = x - residual
+    x = x * residual
+    x = tf.identity(x)
+    o = layers.GlobalAveragePooling2D(keepdims=True)(x)
+    model = keras.Model(inputs=[i], outputs=[o])
+
+    pruned_model = prune.prune_low_magnitude(
+        model,
+        pruning_policy=pruning_policy.PruneForLatencyOnXNNPack(),
+        **self.params,
+    )
+    self.assertEqual(self._count_pruned_layers(pruned_model), 1)
+
 
 if __name__ == '__main__':
   tf.test.main()