code: def testUpdatePruningStepsAndLogsSummaries_CustomTrainingLoop(self):
      log_dir = tempfile.mkdtemp()
      (
          pruned_model,
          loss,
          optimizer,
          x_train,
          y_train,
      ) = self._pruned_model_setup(custom_training_loop=True)
  
      unused_arg = -1
  
      step_callback = pruning_callbacks.UpdatePruningStep()
      log_callback = pruning_callbacks.PruningSummaries(log_dir=log_dir)
      # TODO(tfmot): we need a separate API for custom training loops
      # that doesn't rely on users setting the model and optimizer.
      #
      # Example is currently based on callbacks.py configure_callbacks
      # and model.compile internals.
      step_callback.set_model(pruned_model)
      log_callback.set_model(pruned_model)
      pruned_model.optimizer = optimizer
  
      step_callback.on_train_begin()
      for _ in range(3):
          log_callback.on_epoch_begin(epoch=unused_arg)
          # only one batch given batch_size = 20 and input shape.
          step_callback.on_train_batch_begin(batch=unused_arg)
          inp = np.reshape(
              x_train, [self._BATCH_SIZE, 10]
          )  # original shape: from [10].
          with tf.GradientTape() as tape:
              logits = pruned_model(inp, training=True)
              loss_value = loss(y_train, logits)
              grads = tape.gradient(
                  loss_value, pruned_model.trainable_variables
              )
              optimizer.apply_gradients(
                  zip(grads, pruned_model.trainable_variables)
              )
          step_callback.on_epoch_end(batch=unused_arg)
  
      self.assertEqual(
          3,
          tf.keras.backend.get_value(
              pruned_model.layers[0].pruning_step
          ),
      )
      self.assertEqual(
          3,
          tf.keras.backend.get_value(
              pruned_model.layers[1].pruning_step
          ),
      )
      self._assertLogsExist(log_dir)
  
file_path: model-optimization/tensorflow_model_optimization/python/core/sparsity/keras/pruning_callbacks_test.py
imports: import os
  import tempfile
  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
label_desc: Use image node_id 4 for calling the PruneCallbacksTest obj's underlying member method code with example usage: obj.testUpdatePruningStepsAndLogsSummaries_CustomTrainingLoop() without return types