class_docstr: Wrapper for TFX-BSL bulk inference implementation.
class_docstr_tok: ["Wrapper","for","TFX-BSL","bulk","inference","implementation","."]
code: def expand(
      self, pcoll: beam.PCollection[KeyAndOutput]
  ) -> beam.PCollection[KeyAndOutputMap]:
      # TODO(b/241022420): Set load_override_fn here to avoid loading the model
      # twice.
      return (
          pcoll
          | "TfxBslBulkInference"
          >> run_inference.RunInferencePerModel(
              inference_spec_types=self._aligned_inference_specs
          )
          | "CreateModelNameToPredictionLog"
          >> beam.MapTuple(
              lambda extracts, logs: (  # pylint: disable=g-long-lambda
                  extracts,
                  dict(zip(self._aligned_model_names, logs)),
              )
          )
      )
  
file_path: model-analysis/tensorflow_model_analysis/extractors/tfx_bsl_predictions_extractor.py
imports: import copy
  from typing import Dict, Iterable, List, Optional, Tuple, TypeVar, Union
  import apache_beam
  import tensorflow
  from tensorflow_model_analysis import constants
  from tensorflow_model_analysis.api import types
  from tensorflow_model_analysis.extractors import extractor
  from tensorflow_model_analysis.extractors import inference_base
  from tensorflow_model_analysis.extractors import predictions_extractor
  from tensorflow_model_analysis.proto import config_pb2
  from tensorflow_model_analysis.utils import model_util
  from tfx_bsl.public.beam import run_inference
  from tfx_bsl.public.proto import model_spec_pb2
  from tensorflow_serving.apis import prediction_log_pb2
label_desc: Use image node_id 2 for calling the TfxBslInferenceWrapper obj's underlying member method code with example usage: obj.expand(pcoll) and returns: unknown