ValueError: `prune_low_magnitude` can only prune an object of the following types: keras.models.Sequential, keras functional model, keras.layers.Layer, list of keras.layers.Layer. You passed an object of type: Conv2D.

Start with the prune_model function and the prune_low_magnitude call shown in mbv2_prune.py. Reproduce the failure using MobileNetV2 with TensorFlow 2.16.2, TensorFlow Model Optimization 0.8.0, and Python 3.10.16, then inspect why Conv2D is rejected despite inheriting from Layer. Done means the documented reproduction no longer raises this ValueError and the affected layers can be pruned.

Describe the bug
I am trying to prune MobileNetV2 model using prune_low_magnitude but am running into the following error.

2025-01-19 17:41:03.439078: I tensorflow/core/platform/cpu_feature_guard.cc:210] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
Loading MobileNetV2 model...
Pruning the model...
Traceback (most recent call last):
  File "/Users/esp/Projects/gen_models/mbv2_prune.py", line 89, in <module>
    full_workflow()
  File "/Users/esp/Projects/gen_models/mbv2_prune.py", line 81, in full_workflow
    pruned_model = prune_model(model)
  File "/Users/esp/Projects/gen_models/mbv2_prune.py", line 27, in prune_model
    pruned_layer = prune_low_magnitude(tf.keras.layers.Conv2D(
  File "/Users/esp/Projects/gen_models/venv/lib/python3.10/site-packages/tensorflow_model_optimization/python/core/keras/metrics.py", line 74, in inner
    raise error
  File "/Users/esp/Projects/gen_models/venv/lib/python3.10/site-packages/tensorflow_model_optimization/python/core/keras/metrics.py", line 69, in inner
    results = func(*args, **kwargs)
  File "/Users/esp/Projects/gen_models/venv/lib/python3.10/site-packages/tensorflow_model_optimization/python/core/sparsity/keras/prune.py", line 216, in prune_low_magnitude
    raise ValueError(
ValueError: `prune_low_magnitude` can only prune an object of the following types: keras.models.Sequential, keras functional model, keras.layers.Layer, list of keras.layers.Layer. You passed an object of type: Conv2D.

System information

TensorFlow version (installed from source or binary): 2.16.2

TensorFlow Model Optimization version (installed from source or binary): 0.8.0

Python version: 3.10.16

Describe the expected behavior
The layers should be pruned as Conv2D inherits Layer which as the error says is supported.

Describe the current behavior
Even though Conv2D inherits Layer which as the error says is supported, the code fails.

Code to reproduce the issue

def load_model_mobilenetv2():
    model = MobileNetV2(weights="imagenet", include_top=True, alpha=0.35, input_shape=(224, 224, 3))
    return model

def prune_model(base_model):
    pruning_schedule = PolynomialDecay(
        initial_sparsity=0.0,
        final_sparsity=0.5,
        begin_step=0,
        end_step=1000
    )
    
    pruned_model = Sequential()
    for layer in base_model.layers:
        if isinstance(layer, tf.keras.layers.Dense):
            pruned_layer = prune_low_magnitude(layer, pruning_schedule)
            pruned_model.add(pruned_layer)
        elif isinstance(layer, tf.keras.layers.Conv2D):
            pruned_layer = prune_low_magnitude(tf.keras.layers.Conv2D(
                layer.filters,
                layer.kernel_size,
                layer.strides,
                layer.padding,
                layer.data_format,
                layer.dilation_rate,
                layer.groups,
                layer.activation,
                layer.use_bias,
                layer.kernel_initializer,
                layer.bias_initializer,
                layer.kernel_regularizer,
                layer.bias_regularizer,
                layer.activity_regularizer,
                layer.kernel_constraint,
                layer.bias_constraint,
            ), pruning_schedule)
            pruned_model.add(pruned_layer)
        else:
            pruned_model.add(layer)
    
    # This doesn't work too
    # for layer in base_model.layers:
    #     if isinstance(layer, tf.keras.layers.Dense) or isinstance(layer, tf.keras.layers.Conv2D):
    #         pruned_layer = prune_low_magnitude(layer, pruning_schedule)
    #         pruned_model.add(pruned_layer)
    #     else:
    #         pruned_model.add(layer)

    pruned_model.build(input_shape=(None, 224, 224, 3))
    pruned_model.compile(
        optimizer=tf.keras.optimizers.Adam(),
        loss="categorical_crossentropy",
        metrics=["accuracy"]
    )
    return pruned_model