content: import datetime
  import logging
  from functools import partial, wraps
  from typing import Any, Callable, Optional, Sequence
  
  import torch
  from torch import nn
  from torch.fx.passes.pass_manager import PassManager, inplace_wrapper
  from torch.fx.passes.shape_prop import ShapeProp
  from torch.fx.passes.splitter_base import SplitResult, generate_inputs_for_submodules
  from torch_tensorrt.fx.passes.pass_utils import apply_bfloat_float_conversion
  from torch_tensorrt.fx.utils import LowerPrecision
  
  from ..input_tensor_spec import generate_input_specs
  from ..lower_setting import LowerSetting
  from ..observer import Observer
  from ..passes.remove_duplicate_output_args import remove_duplicate_output_args
  from .graph_opts import common_subexpression_elimination
  from .lower_basic_pass import (  # noqa
      fix_clamp_numerical_limits_to_fp16,
      fix_reshape_batch_dim,
      replace_mutable_op,
      replace_op_with_indices,
      run_const_fold,
  )
  
  _LOGGER: logging.Logger = logging.getLogger(__name__)
  
  
  Input = Sequence[Any]
  
  
  # ----------------------------------------------------------------------
  # OBSERVERS
  # ----------------------------------------------------------------------
  # List of observers. We can subscribe to them by calling its `add(callback)`
  # function from anywhere in code:
  #
  # >>> from torch_tensorrt.fx.lower import FUSE_PASSES_POST_OBSERVER
  # >>> with FUSE_PASSES_POST_OBSERVER.add(print_module_and_input):
  # >>>     # print_module_and_input will be called right after the fuse passes
  # >>>     lower(module, sample_input)
  
  # Observer for the model after the fuse passes.
  FUSE_PASSES_POST_OBSERVER: Observer[Callable[[nn.Module, Input], None]] = Observer(
      "FUSE_PASSES_POST_OBSERVER"
  )
  
  # Observer for the TRT split submodules before lowering
  LOWER_SPLIT_PRE_OBSERVER: Observer[Callable[[str, nn.Module, Input], None]] = Observer(
      "LOWER_SPLIT_PRE_OBSERVER"
  )
  
  # Observer for the TRT split submodules after lowering
  LOWER_SPLIT_POST_OBSERVER: Observer[Callable[[str, nn.Module, Input], None]] = Observer(
      "LOWER_SPLIT_POST_OBSERVER"
  )
  # ----------------------------------------------------------------------
  
  
  def wrapper(fn: Callable, input) -> Callable:
      @wraps(fn)
      def wrapped_fn(gm):
          if isinstance(gm, torch.fx.GraphModule):
              ShapeProp(gm).propagate(*input)
          return fn(gm, input)
  
      return wrapped_fn
  
  
  class LowerPassManagerBuilder:
      """
      Build PassManager for lowering.
  
       Attributes:
          lower_setting: Setting that will be used during process of lowering, see lower_setting.py for the details.
          _trace_func: fx trace function for TRT conversion.
          _split_func: the fx2trt split function.
          _lower_func: function to create and run `TRTInterpreter` to convert `fx.GraphModule`
              into a TensorRT engine.
  
      """
  
      def __init__(
          self,
          lower_setting: LowerSetting,
          trace_func: Callable,
          split_func: Callable,
          lower_func: Callable,
      ):
          self.lower_setting = lower_setting
          self._trace_func = trace_func
          self._split_func = split_func
          self._lower_func = lower_func
  
      def _const_fold_pass(self) -> PassManager:
          passes = [
              wrapper(self._trace_func, self._input),
              run_const_fold,
          ]
          return PassManager.build_from_passlist(passes)
  
      def graph_optimization_pass(self) -> PassManager:
          passes = [
              wrapper(self._trace_func, self._input),
          ]
          for p in self.lower_setting.customized_fuse_pass.passes:
              passes.append(wrapper(p, self._input))
          for p in self.lower_setting.lower_basic_fuse_pass.passes:
              passes.append(wrapper(p, self._input))
          if (
              hasattr(self.lower_setting, "lower_precision")
              and self.lower_setting.lower_precision is LowerPrecision.FP16
          ) or (
              hasattr(self.lower_setting, "precision")
              and self.lower_setting.precision is LowerPrecision.FP16
          ):
              passes.append(wrapper(fix_clamp_numerical_limits_to_fp16, self._input))
  
          passes.append(inplace_wrapper(common_subexpression_elimination))
          passes.append(
              inplace_wrapper(lambda m: FUSE_PASSES_POST_OBSERVER.observe(m, self._input))
          )
          passes.append(fix_reshape_batch_dim)
  
          return PassManager.build_from_passlist(passes)
  
      def graph_optimization_pass_aten(self) -> PassManager:
          passes = []
  
          for p in self.lower_setting.customized_fuse_pass.passes:
              passes.append(wrapper(p, self._input))
          for p in self.lower_setting.lower_basic_fuse_pass.passes:
              passes.append(wrapper(p, self._input))
          # TODO fix this pass for aten graph
          # if (
          #     hasattr(self.lower_setting, "lower_precision")
          #     and self.lower_setting.lower_precision is LowerPrecision.FP16
          # ) or (
          #     hasattr(self.lower_setting, "precision")
          #     and self.lower_setting.precision is LowerPrecision.FP16
          # ):
          #     passes.append(wrapper(fix_clamp_numerical_limits_to_fp16, self._input))
  
          passes.append(
              inplace_wrapper(lambda m: FUSE_PASSES_POST_OBSERVER.observe(m, self._input))
          )
          # TODO we most likely do not need it for aten
          # passes.append(fix_reshape_batch_dim)
  
          return PassManager.build_from_passlist(passes)
  
      def _split_pass(self) -> PassManager:
          passes = [
              partial(
                  self._split_func, inputs=self._input, lower_setting=self.lower_setting
              )
          ]
          passes.append(
              inplace_wrapper(
                  lambda split_result: remove_duplicate_output_args(
                      split_result.split_module, split_result.submodule_inputs.keys()
                  )
              )
          )
          return PassManager.build_from_passlist(passes)
  
      def _trt_lower_pass(self) -> PassManager:
          def lower_func(split_result: SplitResult) -> nn.Module:
              if (
                  hasattr(self.lower_setting, "explicit_batch_dimension")
                  and self.lower_setting.explicit_batch_dimension
                  and self._additional_input
              ):
                  additional_submodule_inputs = generate_inputs_for_submodules(
                      split_result.split_module,
                      self._additional_input,
                      list(split_result.submodule_inputs.keys()),
                  )
              else:
                  additional_submodule_inputs = None
  
              for submod_name, submod_inputs in split_result.submodule_inputs.items():
                  submod = getattr(split_result.split_module, submod_name)
  
                  LOWER_SPLIT_PRE_OBSERVER.observe(submod_name, submod, submod_inputs)
  
                  # Only acc submodules will be lowered.
                  if not submod_name.startswith(split_result.non_acc_submodule_prefix):
                      _LOGGER.info(f"Now lowering submodule {submod_name}")
                      lowering_start_time = datetime.datetime.now()
  
                      self.lower_setting.input_specs = generate_input_specs(
                          submod_inputs,
                          self.lower_setting,
                          (
                              additional_submodule_inputs[submod_name]
                              if additional_submodule_inputs
                              else None
                          ),
                      )
                      lowered_module = self._lower_func(
                          submod, submod_inputs, self.lower_setting, submod_name
                      )
                      setattr(split_result.split_module, submod_name, lowered_module)
                      LOWER_SPLIT_POST_OBSERVER.observe(
                          submod_name, lowered_module, submod_inputs
                      )
                      _LOGGER.info(
                          f"Lowering submodule {submod_name} elapsed time {datetime.datetime.now() - lowering_start_time}"
                      )
  
              return split_result.split_module
  
          return PassManager.build_from_passlist([lower_func])
  
      def _default_lower_pass(self) -> PassManager:
          def lower_func(split_result: SplitResult) -> nn.Module:
              if self._additional_input:
                  additional_submodule_inputs = generate_inputs_for_submodules(
                      split_result.split_module,
                      self._additional_input,
                      list(split_result.submodule_inputs.keys()),
                  )
              else:
                  additional_submodule_inputs = None
  
              for submod_name, submod_inputs in split_result.submodule_inputs.items():
                  submod = getattr(split_result.split_module, submod_name)
  
                  LOWER_SPLIT_PRE_OBSERVER.observe(submod_name, submod, submod_inputs)
                  # Only acc submodules will be lowered.
                  if not submod_name.startswith(split_result.non_acc_submodule_prefix):
                      _LOGGER.info(f"ACC submodule graph: {submod.graph}")
                      lowering_start_time = datetime.datetime.now()
  
                      self.lower_setting.additional_inputs = (
                          (
                              additional_submodule_inputs[submod_name]
                              if additional_submodule_inputs
                              else None
                          ),
                      )
  
                      lowered_module = self._lower_func(
                          submod, submod_inputs, self.lower_setting, submod_name
                      )
                      setattr(split_result.split_module, submod_name, lowered_module)
                      LOWER_SPLIT_POST_OBSERVER.observe(
                          submod_name, lowered_module, submod_inputs
                      )
                      _LOGGER.info(
                          f"Lowering submodule {submod_name} elapsed time {datetime.datetime.now() - lowering_start_time}"
                      )
                  else:
                      _LOGGER.info(f"GPU submodule graph: {submod.graph}")
                      apply_bfloat_float_conversion(submod, submod_inputs, submod_name)
  
              return split_result.split_module
  
          return PassManager.build_from_passlist([lower_func])
  
      def _default_replace_mutable_op_pass(self) -> PassManager:
          return PassManager.build_from_passlist([replace_mutable_op])
  
      def build_trt_lower_pipeline(
          self, input: Input, additional_input: Optional[Input] = None
      ) -> PassManager:
          self._input = input
          self._additional_input = additional_input
          passes = []
  
          passes.append(self._default_replace_mutable_op_pass())
          passes.append(self._const_fold_pass())
          passes.append(self.graph_optimization_pass())
          passes.append(self._split_pass())
          passes.append(self._trt_lower_pass())
  
          pm = PassManager.build_from_passlist(passes)
          return pm
  
      def build_aten2trt_lower_pipeline(
          self, input: Input, additional_input: Optional[Input] = None
      ) -> PassManager:
          self._input = input
          self._additional_input = additional_input
          passes = []
          passes.append(
              wrapper(self._trace_func, self._input),
          )
          passes.append(self.graph_optimization_pass_aten())
          passes.append(self._split_pass())
          passes.append(self._trt_lower_pass())
  
          pm = PassManager.build_from_passlist(passes)
          return pm
  
      def build_default_lower_pipeline(
          self, input: Input, additional_input: Optional[Input] = None
      ) -> PassManager:
          self._input = input
          self._additional_input = additional_input
          passes = []
  
          passes.append(self._default_replace_mutable_op_pass())
          passes.append(self._const_fold_pass())
          passes.append(self.graph_optimization_pass())
          passes.append(self._split_pass())
          passes.append(self._default_lower_pass())
  
          pm = PassManager.build_from_passlist(passes)
          return pm
  