content: from typing import Optional
  
  import tensorrt as trt
  import torch
  from tensorrt import ITensor as TRTTensor
  from torch.fx.node import Target
  from torch_tensorrt import _enums
  from torch_tensorrt.dynamo._SourceIR import SourceIR
  from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
  from torch_tensorrt.dynamo.conversion.converter_utils import (
      broadcast,
      cast_trt_tensor,
      get_trt_tensor,
      set_layer_name,
  )
  
  
  def matrix_multiply(
      ctx: ConversionContext,
      target: Target,
      source_ir: Optional[SourceIR],
      name: str,
      input: TRTTensor,
      other: TRTTensor,
      input_matrix_op: trt.MatrixOperation = trt.MatrixOperation.NONE,
      other_matrix_op: trt.MatrixOperation = trt.MatrixOperation.NONE,
  ) -> TRTTensor:
      if not isinstance(input, trt.ITensor):
          input = get_trt_tensor(ctx, input, f"{name}_input")
      if not isinstance(other, trt.ITensor):
          other = get_trt_tensor(
              ctx,
              other,
              f"{name}_other",
              dtype=_enums.dtype._from(input.dtype).to(torch.dtype),
          )
  
      preset_diff = 0
  
      if len(input.shape) == 1:
          preset_diff -= 1
          input_matrix_op = trt.MatrixOperation.VECTOR
  
      if len(other.shape) == 1:
          preset_diff += 1
          other_matrix_op = trt.MatrixOperation.VECTOR
  
      input, other = broadcast(
          ctx, input, other, f"{name}_input", f"{name}_other", preset_diff
      )
      # Get the original input dtype
      input_dtype = _enums.dtype._from(input.dtype).to(torch.dtype)
  
      if (
          ctx.net.get_flag(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)
          and ctx.compilation_settings.use_fp32_acc
          and input_dtype == torch.float16
      ):
          input = cast_trt_tensor(ctx, input, torch.float32, f"{name}_input_casted")
          other = cast_trt_tensor(ctx, other, torch.float32, f"{name}_other_casted")
  
      matmul_layer = ctx.net.add_matrix_multiply(
          input, input_matrix_op, other, other_matrix_op
      )
      matmul_output = matmul_layer.get_output(0)
  
      if (
          ctx.net.get_flag(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)
          and ctx.compilation_settings.use_fp32_acc
          and input_dtype == torch.float16
      ):
          matmul_output = cast_trt_tensor(
              ctx, matmul_output, input_dtype, f"{name}_output_casted"
          )
  
      set_layer_name(matmul_layer, target, name, source_ir)
      return matmul_output
  