code: def build_layerwise_model(input_shape, **pruning_params):
      return tf.keras.Sequential(
          [
              prune.prune_low_magnitude(
                  l.Conv2D(
                      32,
                      5,
                      padding="same",
                      activation="relu",
                      input_shape=input_shape,
                  ),
                  **pruning_params,
              ),
              l.MaxPooling2D((2, 2), (2, 2), padding="same"),
              prune.prune_low_magnitude(
                  l.Conv2D(64, 5, padding="same"), **pruning_params
              ),
              l.BatchNormalization(),
              l.ReLU(),
              l.MaxPooling2D((2, 2), (2, 2), padding="same"),
              l.Flatten(),
              prune.prune_low_magnitude(
                  l.Dense(1024, activation="relu"), **pruning_params
              ),
              l.Dropout(0.4),
              l.Dense(num_classes, activation="softmax"),
          ]
      )
  
file_path: model-optimization/tensorflow_model_optimization/python/examples/sparsity/keras/mnist/mnist_e2e_sparsity2x4.py
imports: from __future__ import print_function
  from absl import app
  import tensorflow
  from tensorflow_model_optimization.python.core.keras import test_utils
  from tensorflow_model_optimization.python.core.sparsity.keras import prune
  from tensorflow_model_optimization.python.core.sparsity.keras import pruning_callbacks
  from tensorflow_model_optimization.python.core.sparsity.keras import pruning_schedule
  from tensorflow_model_optimization.python.core.sparsity.keras import pruning_utils
  from tensorflow_model_optimization.python.core.sparsity.keras import pruning_wrapper
label_desc: Use image node_id 2 for calling a global function with example usage: build_layerwise_model(input_shape) and returns: tf