code: def testRekeyPredictionsInFeaturesForPrematerializedPredictions(self):
      model_spec1 = config_pb2.ModelSpec(
          name="model1", prediction_key="prediction"
      )
      model_spec2 = config_pb2.ModelSpec(
          name="model2",
          prediction_keys={
              "output1": "prediction1",
              "output2": "prediction2",
          },
      )
      eval_config = config_pb2.EvalConfig(
          model_specs=[model_spec1, model_spec2]
      )
      schema = text_format.Parse(
          """
          tensor_representation_group {
            key: ""
            value {
              tensor_representation {
                key: "fixed_int"
                value {
                  dense_tensor {
                    column_name: "fixed_int"
                  }
                }
              }
            }
          }
          feature {
            name: "prediction"
            type: FLOAT
            shape: { }
            presence: { min_fraction: 1 }
          }
          feature {
            name: "prediction1"
            type: FLOAT
            shape: { }
            presence: { min_fraction: 1 }
          }
          feature {
            name: "prediction2"
            type: FLOAT
            shape: { }
            presence: { min_fraction: 1 }
          }
          feature {
            name: "fixed_int"
            type: INT
          }
          """,
          schema_pb2.Schema(),
      )
      # TODO(b/73109633): Remove when field is removed or its default changes to
      # False.
      if hasattr(schema, "generate_legacy_feature_spec"):
          schema.generate_legacy_feature_spec = False
      (
          tfx_io,
          feature_extractor,
      ) = self._create_tfxio_and_feature_extractor(eval_config, schema)
  
      examples = [
          self._makeExample(
              prediction=1.0,
              prediction1=1.0,
              prediction2=0.0,
              fixed_int=1,
          ),
          self._makeExample(
              prediction=1.0,
              prediction1=1.0,
              prediction2=1.0,
              fixed_int=1,
          ),
      ]
      num_examples = len(examples)
  
      prediction_extractor = predictions_extractor.PredictionsExtractor(
          eval_config=eval_config, eval_shared_model=None
      )
  
      with beam.Pipeline() as pipeline:
          # pylint: disable=no-value-for-parameter
          result = (
              pipeline
              | "Create"
              >> beam.Create(
                  [e.SerializeToString() for e in examples],
                  reshuffle=False,
              )
              | "BatchExamples"
              >> tfx_io.BeamSource(batch_size=num_examples)
              | "InputsToExtracts"
              >> model_eval_lib.BatchedInputsToExtracts()
              | feature_extractor.stage_name
              >> feature_extractor.ptransform
              | prediction_extractor.stage_name
              >> prediction_extractor.ptransform
          )
          # pylint: enable=no-value-for-parameter
  
          def check_result(got):
              try:
                  self.assertLen(got, 1)
                  for model_name in ("model1", "model2"):
                      self.assertIn(
                          model_name, got[0][constants.PREDICTIONS_KEY]
                      )
                  self.assertAllClose(
                      np.array([1.0, 1.0]),
                      got[0][constants.PREDICTIONS_KEY]["model1"],
                  )
  
                  self.assertAllClose(
                      {
                          "output1": np.array([1.0, 1.0]),
                          "output2": np.array([0.0, 1.0]),
                      },
                      got[0][constants.PREDICTIONS_KEY]["model2"],
                  )
  
              except AssertionError as err:
                  raise util.BeamAssertException(err)
  
          util.assert_that(result, check_result)
  
imports: import os
  from absl.testing import parameterized
  import apache_beam
  from apache_beam.testing import util
  import numpy
  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 dnn_classifier
  from tensorflow_model_analysis.eval_saved_model.example_trainers import fixed_prediction_estimator_extra_fields
  from tensorflow_model_analysis.eval_saved_model.example_trainers import multi_head
  from tensorflow_model_analysis.extractors import features_extractor
  from tensorflow_model_analysis.extractors import predictions_extractor
  from tensorflow_model_analysis.proto import config_pb2
  from tfx_bsl.tfxio import tensor_adapter
  from tfx_bsl.tfxio import test_util
  from google.protobuf import text_format
  from tensorflow_metadata.proto.v0 import schema_pb2
label_desc: Use image node_id 13 for calling the PredictionsExtractorTest obj's underlying member method code with example usage: obj.testRekeyPredictionsInFeaturesForPrematerializedPredictions() without return types