code: def build_functional_model(input_shape):
      inp = tf.keras.Input(shape=input_shape)
      x = l.Conv2D(32, 5, padding="same", activation="relu")(inp)
      x = l.MaxPooling2D((2, 2), (2, 2), padding="same")(x)
      x = l.BatchNormalization()(x)
      x = l.Conv2D(64, 5, padding="same", activation="relu")(x)
      x = l.MaxPooling2D((2, 2), (2, 2), padding="same")(x)
      x = l.Flatten()(x)
      x = l.Dense(1024, activation="relu")(x)
      x = l.Dropout(0.4)(x)
      out = l.Dense(num_classes, activation="softmax")(x)
  
      return tf.keras.models.Model([inp], [out])
  
file_path: model-optimization/tensorflow_model_optimization/python/examples/sparsity/keras/mnist/mnist_cnn.py
imports: from __future__ import print_function
  from absl import app
  from absl import flags
  import tensorflow
  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
label_desc: Use image node_id 2 for calling a global function with example usage: build_functional_model(input_shape) and returns: tf