code: def train(model, x_train, y_train, x_test, y_test):
      model.compile(
          loss=tf.keras.losses.categorical_crossentropy,
          optimizer="adam",
          metrics=["accuracy"],
      )
      model.run_eagerly = True
  
      # Print the model summary.
      model.summary()
  
      # Add a pruning step callback to peg the pruning step to the optimizer's
      # step. Also add a callback to add pruning summaries to tensorboard
      callbacks = [
          pruning_callbacks.UpdatePruningStep(),
          pruning_callbacks.PruningSummaries(log_dir="/tmp/logs"),
      ]
  
      model.fit(
          x_train,
          y_train,
          batch_size=batch_size,
          epochs=epochs,
          verbose=1,
          callbacks=callbacks,
          validation_data=(x_test, y_test),
      )
      score = model.evaluate(x_test, y_test, verbose=0)
      print("Test loss:", score[0])
      print("Test accuracy:", score[1])
  
      # Check sparsity 2x4 type before stripping pruning
      is_pruned_2x4 = check_model_sparsity_2x4(model)
      print("Pass the check for sparsity 2x4: ", is_pruned_2x4)
  
      model = prune.strip_pruning(model)
      return model
  
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 3 for calling a global function with example usage: train(model, x_train, y_train, x_test, y_test) and returns: model