class_docstr: A container that holds info about partial batches for `predict()`.
docstr: Removes prediction output that corresponds to padded input.
code: def apply_mask(self, prediction_result):
      """Removes prediction output that corresponds to padded input."""
      padding_mask = backend.get_value(self.padding_mask)
      assert len(padding_mask.shape) == 1
  
      if len(self.output_shape) == 1:
          prediction = np.take(
              prediction_result,
              np.nonzero(padding_mask[: len(prediction_result)]),
              axis=0,
          )
          if prediction.shape[0] == 1:
              prediction = np.squeeze(prediction, axis=0)
          return prediction
  
      else:
          predictions = []
          for i in range(len(self.output_shape)):
              prediction = prediction_result[i]
              prediction = np.take(
                  prediction,
                  np.nonzero(padding_mask[: len(prediction)]),
                  axis=0,
              )
              predictions.append(np.squeeze(prediction))
  
          return predictions
  
file_path: tensorflow-upstream/tensorflow/python/keras/engine/partial_batch_padding_handler.py
imports: import numpy
  from tensorflow.python.framework import tensor_util
  from tensorflow.python.keras import backend
  from tensorflow.python.ops import array_ops
  from tensorflow.python.util import nest
label_desc: Use image node_id 5 for calling the PartialBatchPaddingHandler obj's underlying member method code with example usage: obj.apply_mask(prediction_result) and returns: prediction, predictions