code: def testPredictMultipleExampleRefPerRawExampleBytes(self):
      temp_eval_export_dir = self._getEvalExportDir()
      (
          _,
          eval_export_dir,
      ) = fake_multi_examples_per_input_estimator.fake_multi_examples_per_input_estimator(
          None, temp_eval_export_dir
      )
      eval_shared_model = model_eval_lib.default_eval_shared_model(
          eval_saved_model_path=eval_export_dir
      )
  
      # The trailing zeros make an "empty" output batch.
      raw_example_bytes = ["0", "3", "1", "0", "2", "0", "0", "0", "0"]
  
      def check_result(got):
          try:
              self.assertLen(got, 6)
              self.assertEqual(
                  ["3", "3", "3", "1", "2", "2"],
                  [extracts[constants.INPUT_KEY] for extracts in got],
              )
  
              for item in got:
                  self.assertIn(
                      constants.FEATURES_PREDICTIONS_LABELS_KEY, item
                  )
                  fpl = item[constants.FEATURES_PREDICTIONS_LABELS_KEY]
                  self.assertIn("input_index", fpl.features)
                  self.assertIn("example_count", fpl.features)
                  self.assertIn("intra_input_index", fpl.features)
  
          except AssertionError as err:
              raise util.BeamAssertException(err)
  
      with beam.Pipeline() as pipeline:
          predict_extracts = (
              pipeline
              | beam.Create(raw_example_bytes, reshuffle=False)
              # Our diagnostic outputs, pass types.Extracts throughout, however our
              # aggregating functions do not use this interface.
              | beam.Map(lambda x: {constants.INPUT_KEY: x})
              | "Predict"
              >> predict_extractor._TFMAPredict(
                  eval_shared_models={"": eval_shared_model},
                  desired_batch_size=3,
              )
          )
  
          util.assert_that(predict_extracts, check_result)
  
imports: import os
  from absl.testing import parameterized
  import apache_beam
  from apache_beam.testing import util
  import tensorflow
  from tensorflow_model_analysis import constants
  from tensorflow_model_analysis.api import model_eval_lib
  from tensorflow_model_analysis.eval_saved_model import testutil
  from tensorflow_model_analysis.eval_saved_model.example_trainers import batch_size_limited_classifier
  from tensorflow_model_analysis.eval_saved_model.example_trainers import fake_multi_examples_per_input_estimator
  from tensorflow_model_analysis.eval_saved_model.example_trainers import linear_classifier
  from tensorflow_model_analysis.extractors import legacy_predict_extractor
  from tensorflow_model_analysis.proto import config_pb2
  from tfx_bsl.tfxio import raw_tf_record
label_desc: Use image node_id 5 for calling the PredictExtractorTest obj's underlying member method code with example usage: obj.testPredictMultipleExampleRefPerRawExampleBytes() without return types