content: # mypy: disallow-untyped-decorators=False
  
  import logging
  from typing import Dict, Sequence, Tuple, Union
  
  import tensorrt as trt
  from torch.fx.node import Argument, Target
  from torch_tensorrt._features import needs_trtllm_for_nccl
  from torch_tensorrt.dynamo._SourceIR import SourceIR
  from torch_tensorrt.dynamo.conversion import impl
  from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
  from torch_tensorrt.dynamo.conversion._ConverterRegistry import (
      dynamo_tensorrt_converter,
  )
  from torch_tensorrt.dynamo.lowering.passes.fuse_distributed_ops import (
      tensorrt_fused_nccl_all_gather_op,
      tensorrt_fused_nccl_reduce_scatter_op,
  )
  
  _LOGGER: logging.Logger = logging.getLogger(__name__)
  
  
  @needs_trtllm_for_nccl
  @dynamo_tensorrt_converter(tensorrt_fused_nccl_all_gather_op)
  def fused_nccl_gather(
      ctx: ConversionContext,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[trt.ITensor, Sequence[trt.ITensor]]:
      return impl.nccl_ops.nccl_gather(
          ctx,
          target,
          SourceIR.ATEN,
          name,
          [args[0]],
      )
  
  
  @needs_trtllm_for_nccl
  @dynamo_tensorrt_converter(tensorrt_fused_nccl_reduce_scatter_op)
  def fused_nccl_reduce_scatter(
      ctx: ConversionContext,
      target: Target,
      args: Tuple[Argument, ...],
      kwargs: Dict[str, Argument],
      name: str,
  ) -> Union[trt.ITensor, Sequence[trt.ITensor]]:
      return impl.nccl_ops.nccl_reduce_scatter(
          ctx,
          target,
          SourceIR.ATEN,
          name,
          [args[0]],
      )
  