code: def testBatchSizeLimitWithKerasModel(self):
      input1 = tf.keras.layers.Input(
          shape=(1,), batch_size=1, name="input1"
      )
      input2 = tf.keras.layers.Input(
          shape=(1,), batch_size=1, name="input2"
      )
  
      inputs = [input1, input2]
      input_layer = tf.keras.layers.concatenate(inputs)
  
      def add_1(tensor):
          return tf.add_n([tensor, tf.constant(1.0, shape=(1, 2))])
  
      assert_layer = tf.keras.layers.Lambda(add_1)(input_layer)
  
      model = tf.keras.models.Model(inputs, assert_layer)
      model.compile(
          optimizer=tf.keras.optimizers.Adam(lr=0.001),
          loss=tf.keras.losses.binary_crossentropy,
          metrics=["accuracy"],
      )
  
      export_dir = self._getExportDir()
      model.save(export_dir, save_format="tf")
  
      eval_config = config_pb2.EvalConfig(
          model_specs=[config_pb2.ModelSpec()]
      )
      eval_shared_model = self.createTestEvalSharedModel(
          eval_saved_model_path=export_dir,
          tags=[tf.saved_model.SERVING],
      )
      (
          tfx_io,
          feature_extractor,
      ) = self._create_tfxio_and_feature_extractor(
          eval_config,
          text_format.Parse(
              """
          tensor_representation_group {
            key: ""
            value {
              tensor_representation {
                key: "input1"
                value {
                  dense_tensor {
                    column_name: "input1"
                    shape { dim { size: 1 } }
                  }
                }
              }
              tensor_representation {
                key: "input2"
                value {
                  dense_tensor {
                    column_name: "input2"
                    shape { dim { size: 1 } }
                  }
                }
              }
            }
          }
          feature {
            name: "input1"
            type: FLOAT
          }
          feature {
            name: "input2"
            type: FLOAT
          }
          """,
              schema_pb2.Schema(),
          ),
      )
      prediction_extractor = predictions_extractor.PredictionsExtractor(
          eval_config=eval_config, eval_shared_model=eval_shared_model
      )
  
      examples = []
      for _ in range(4):
          examples.append(self._makeExample(input1=0.0, input2=1.0))
  
      with beam.Pipeline() as pipeline:
          predict_extracts = (
              pipeline
              | "Create"
              >> beam.Create(
                  [e.SerializeToString() for e in examples],
                  reshuffle=False,
              )
              | "BatchExamples" >> tfx_io.BeamSource(batch_size=1)
              | "InputsToExtracts"
              >> model_eval_lib.BatchedInputsToExtracts()
              | feature_extractor.stage_name
              >> feature_extractor.ptransform
              | prediction_extractor.stage_name
              >> prediction_extractor.ptransform
          )
  
          def check_result(got):
              try:
                  self.assertLen(got, 4)
                  # We can't verify the actual predictions, but we can verify the keys.
                  for item in got:
                      self.assertIn(constants.PREDICTIONS_KEY, item)
  
              except AssertionError as err:
                  raise util.BeamAssertException(err)
  
          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 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 11 for calling the PredictionsExtractorTest obj's underlying member method code with example usage: obj.testBatchSizeLimitWithKerasModel() and returns: tf