code: def testPrunesSimpleDenseModel(self, distribution):
      with distribution.scope():
          model = prune.prune_low_magnitude(
              keras_test_utils.build_simple_dense_model(), **self.params
          )
          model.compile(
              loss="categorical_crossentropy",
              optimizer="sgd",
              metrics=["accuracy"],
          )
  
      # Model hasn't been trained yet. Sparsity 0.0
      test_utils.assert_model_sparsity(self, 0.0, model)
  
      # Simple unpruned model. No sparsity.
      model.fit(
          np.random.rand(20, 10),
          keras.utils.to_categorical(
              np.random.randint(5, size=(20, 1)), 5
          ),
          epochs=2,
          callbacks=[pruning_callbacks.UpdatePruningStep()],
          batch_size=20,
      )
      model.predict(np.random.rand(20, 10))
      test_utils.assert_model_sparsity(self, 0.5, model)
  
      _, keras_file = tempfile.mkstemp(".h5")
      keras.models.save_model(model, keras_file)
  
      with prune.prune_scope():
          loaded_model = keras.models.load_model(keras_file)
  
      test_utils.assert_model_sparsity(self, 0.5, loaded_model)
  
file_path: model-optimization/tensorflow_model_optimization/python/core/sparsity/keras/prune_distributed_test.py
imports: import tempfile
  from absl.testing import parameterized
  import numpy
  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 test_utils
label_desc: Use image node_id 2 for calling the PruneDistributedTest obj's underlying member method code with example usage: obj.testPrunesSimpleDenseModel(distribution) without return types