Fixed comments

Author: cehongwang

py/torch_tensorrt/_compile.py

@@ -243,11 +243,11 @@ def compile(
     elif target_ir == _IRType.dynamo:
         # Prepare torch and torchtrt inputs
         if not arg_inputs and not inputs:
-            raise AssertionError("'arg_input' and 'input' should not both be None.")
+            raise AssertionError("'arg_inputs' and 'inputs' should not both be None.")
 
         elif arg_inputs and inputs:
             raise AssertionError(
-                "'arg_input' and 'input' should not be used at the same time."
+                "'arg_inputs' and 'inputs' should not be used at the same time."
             )
         arg_inputs = inputs or arg_inputs
 
@@ -367,11 +367,11 @@ def convert_method_to_trt_engine(
     elif target_ir == _IRType.dynamo:
         # Prepare torch and torchtrt inputs
         if not arg_inputs and not inputs:
-            raise AssertionError("'arg_input' and 'input' should not both be None.")
+            raise AssertionError("'arg_inputs' and 'inputs' should not both be None.")
 
         elif arg_inputs and inputs:
             raise AssertionError(
-                "'arg_input' and 'input' should not be used at the same time."
+                "'arg_inputs' and 'inputs' should not be used at the same time."
             )
         arg_inputs = arg_inputs or inputs
 
@@ -450,7 +450,7 @@ def save(
     output_format: str = "exported_program",
     inputs: Optional[Sequence[torch.Tensor]] = None,
     arg_inputs: Optional[Sequence[torch.Tensor]] = None,
-    kwargs_inputs: Optional[dict[str, Any]] = None,
+    kwarg_inputs: Optional[dict[str, Any]] = None,
     retrace: bool = False,
 ) -> None:
     """
@@ -475,15 +475,15 @@ def save(
         )
     if arg_inputs and inputs:
         raise AssertionError(
-            "'arg_input' and 'input' should not be used at the same time."
+            "'arg_inputs' and 'inputs' should not be used at the same time."
         )
 
     arg_inputs = inputs or arg_inputs
 
-    if kwargs_inputs is None:
-        kwargs_inputs = {}
+    if kwarg_inputs is None:
+        kwarg_inputs = {}
 
-    if kwargs_inputs and any(value is None for value in kwargs_inputs.values()):
+    if kwarg_inputs and any(value is None for value in kwarg_inputs.values()):
         raise ValueError("kwargs should not include None.")
     if output_format not in accepted_formats:
         raise ValueError(
@@ -518,20 +518,20 @@ def save(
         # The module type is torch.fx.GraphModule
         if output_format == "torchscript":
             module_ts = torch.jit.trace(
-                module, arg_inputs, example_kwarg_inputs=kwargs_inputs
+                module, arg_inputs, example_kwarg_inputs=kwarg_inputs
             )
             torch.jit.save(module_ts, file_path)
         else:
             if not retrace:
                 from torch_tensorrt.dynamo._exporter import export
 
-                exp_program = export(module, arg_inputs, kwargs_inputs)
+                exp_program = export(module, arg_inputs, kwarg_inputs)
                 torch.export.save(exp_program, file_path)
             else:
                 from torch._higher_order_ops.torchbind import enable_torchbind_tracing
 
                 with enable_torchbind_tracing():
                     exp_program = torch.export.export(
-                        module, tuple(arg_inputs), kwargs=kwargs_inputs, strict=False
+                        module, tuple(arg_inputs), kwargs=kwarg_inputs, strict=False
                     )
                     torch.export.save(exp_program, file_path)

py/torch_tensorrt/dynamo/_compiler.py

@@ -187,11 +187,11 @@ def compile(
 
     # Aliasing inputs to arg_inputs for better understanding
     if not arg_inputs and not inputs:
-        raise AssertionError("'arg_input' and 'input' should not both be None.")
+        raise AssertionError("'arg_inputs' and 'inputs' should not both be None.")
 
     elif arg_inputs and inputs:
         raise AssertionError(
-            "'arg_input' and 'input' should not be used at the same time."
+            "'arg_inputs' and 'inputs' should not be used at the same time."
         )
 
     arg_inputs = inputs or arg_inputs
@@ -460,13 +460,13 @@ def contains_metadata(gm: torch.fx.GraphModule) -> bool:
 
             trt_modules[name] = trt_module
 
-    torch_sample_inputs = get_torch_inputs(
+    torch_sample_arg_inputs = get_torch_inputs(
         sample_arg_inputs, to_torch_device(settings.device)
     )
     torch_sample_kwarg_inputs = get_torch_inputs(
         sample_kwarg_inputs, to_torch_device(settings.device)
     )
-    sample_outputs = gm(*torch_sample_inputs, **torch_sample_kwarg_inputs)
+    sample_outputs = gm(*torch_sample_arg_inputs, **torch_sample_kwarg_inputs)
 
     if not isinstance(sample_outputs, (list, tuple)):
         sample_outputs = [sample_outputs]
@@ -609,23 +609,22 @@ def convert_module_to_trt_engine(
             stacklevel=2,
         )
     if not arg_inputs and not inputs:
-        raise AssertionError("'arg_input' and 'input' should not both be None.")
+        raise AssertionError("'arg_inputs' and 'inputs' should not both be None.")
 
     elif arg_inputs and inputs:
         raise AssertionError(
-            "'arg_input' and 'input' should not be used at the same time."
+            "'arg_inputs' and 'inputs' should not be used at the same time."
         )
 
     arg_inputs = inputs or arg_inputs
 
-    arg_input_list = list(arg_inputs) if arg_inputs is not None else []
     torch_executed_ops = torch_executed_ops if torch_executed_ops is not None else set()
     if kwarg_inputs is None:
         kwarg_inputs = {}
     # Prepare torch_trt inputs
-    arg_input_list = prepare_inputs(arg_input_list)
+    arg_inputs = prepare_inputs(arg_inputs)
     kwarg_input_list = prepare_inputs(kwarg_inputs)
-    flattened_input_list = arg_input_list + flatten_dict_value(kwarg_input_list)
+    flattened_input_list = arg_inputs + flatten_dict_value(kwarg_input_list)
     device = to_torch_tensorrt_device(device)
     enabled_precisions = {dtype._from(e) for e in enabled_precisions}
 
@@ -679,7 +678,7 @@ def convert_module_to_trt_engine(
         interpreter_result = interpret_module_to_result(
             gm,
             inputs=flattened_input_list,
-            arg_inputs=arg_input_list,
+            arg_inputs=arg_inputs,
             kwarg_inputs=kwarg_input_list,
             settings=settings,
         )

py/torch_tensorrt/dynamo/_exporter.py

@@ -22,25 +22,25 @@
 def export(
     gm: torch.fx.GraphModule,
     inputs: Sequence[torch.Tensor],
-    kwargs_inputs: Optional[dict[str, Any]] = None,
+    kwarg_inputs: Optional[dict[str, Any]] = None,
 ) -> ExportedProgram:
     """Export the result of TensorRT compilation into the desired output format.
 
     Arguments:
         gm (torch.fx.GraphModule): Compiled Torch-TensorRT module, generated by ``torch_tensorrt.dynamo.compile``
         inputs (torch.Tensor): Torch input tensors
     """
-    if kwargs_inputs is None:
-        kwargs_inputs = {}
-    patched_module = transform(gm, inputs, kwargs_inputs)
+    if kwarg_inputs is None:
+        kwarg_inputs = {}
+    patched_module = transform(gm, inputs, kwarg_inputs)
     exp_program = create_trt_exp_program(patched_module)
     return exp_program
 
 
 def transform(
     gm: torch.fx.GraphModule,
     inputs: Sequence[torch.Tensor],
-    kwargs_inputs: Optional[dict[str, Any]] = None,
+    kwarg_inputs: Optional[dict[str, Any]] = None,
 ) -> torch.fx.GraphModule:
     """
     Transforms the graphmodule by inlining Pytorch and TensorRT submodules.
@@ -55,11 +55,11 @@ def transform(
     """
     # Make a copy the graph since this function transforms the input graph and changes it's attributes.
     # This transformed graph is meant to be consumed by `create_trt_exp_program`
-    if kwargs_inputs is None:
-        kwargs_inputs = {}
+    if kwarg_inputs is None:
+        kwarg_inputs = {}
     gm = copy.deepcopy(gm)
     # Run shape analysis
-    _, outputs_map = partitioning.run_shape_analysis(gm, inputs, kwargs_inputs)
+    _, outputs_map = partitioning.run_shape_analysis(gm, inputs, kwarg_inputs)
 
     # Inline TensorRT submodules
     inline_trt_modules(gm, outputs_map)

py/torch_tensorrt/dynamo/_tracer.py

@@ -59,11 +59,11 @@ def trace(
 
     # Set log level at the top of compilation (torch_tensorrt.dynamo)
     if not arg_inputs and not inputs:
-        raise AssertionError("'arg_input' and 'input' should not both be None.")
+        raise AssertionError("'arg_inputs' and 'inputs' should not both be None.")
 
     elif arg_inputs and inputs:
         raise AssertionError(
-            "'arg_input' and 'input' should not be used at the same time."
+            "'arg_inputs' and 'inputs' should not be used at the same time."
         )
     arg_inputs = inputs or arg_inputs
 
@@ -79,7 +79,6 @@ def trace(
     torch_kwarg_inputs = get_torch_inputs(kwarg_inputs, device)
     dynamic_shapes = get_dynamic_shapes_args(mod, arg_inputs)
     dynamic_shapes.update(get_dynamic_shapes_kwargs(kwarg_inputs))
-    # breakpoint()
     exp_program = export(
         mod,
         tuple(torch_arg_inputs),

py/torch_tensorrt/dynamo/backend/backends.py

@@ -84,7 +84,7 @@ def _pretraced_backend(
             ]
 
             # Remove detach nodes
-            remove_detach(gm, torch_inputs)
+            remove_detach(gm)
 
             # Invoke AOTAutograd to translate operators to aten
             gm = aot_export_joint_simple(

py/torch_tensorrt/dynamo/partitioning/common.py

@@ -131,7 +131,7 @@ def construct_submodule_inputs(module: torch.fx.GraphModule) -> Sequence[Input]:
 def run_shape_analysis(
     parent_module: torch.fx.GraphModule,
     inputs: Sequence[Input],
-    kwargs_inputs: Optional[dict[str, Any]] = None,
+    kwarg_inputs: Optional[dict[str, Any]] = None,
 ) -> Tuple[Dict[Any, Sequence[Any]], Dict[Any, Sequence[Any]]]:
     submod_inputs_shape_map: Dict[Any, Sequence[Any]] = {}
     submod_outputs_shape_map: Dict[Any, Sequence[Any]] = {}
@@ -147,13 +147,13 @@ def get_submodule_io(
         sub_outputs = outputs
         return
 
-    if kwargs_inputs is None:
-        kwargs_inputs = {}
+    if kwarg_inputs is None:
+        kwarg_inputs = {}
     # Iterate through submodules (both Torch and TRT) and store IO shapes
     for name, _ in parent_module.named_children():
         submodule = getattr(parent_module, name)
         handle = submodule.register_forward_hook(get_submodule_io)
-        parent_module(*inputs, **kwargs_inputs)
+        parent_module(*inputs, **kwarg_inputs)
         handle.remove()
         submod_inputs_shape_map[name] = (
             [input.shape for input in sub_inputs]

tests/py/dynamo/models/test_export_kwargs_serde.py

@@ -74,7 +74,7 @@ def forward(self, x, b=5, c=None, d=None):
 
     # Save the module
     trt_ep_path = os.path.join(tempfile.gettempdir(), "compiled.ep")
-    torchtrt.save(trt_gm, trt_ep_path, inputs=args, kwargs_inputs=kwargs)
+    torchtrt.save(trt_gm, trt_ep_path, inputs=args, kwarg_inputs=kwargs)
     # Clean up model env
     torch._dynamo.reset()
 
@@ -133,7 +133,7 @@ def forward(self, x, b=5, c=None, d=None):
 
     # Save the module
     trt_ep_path = os.path.join(tempfile.gettempdir(), "compiled.ep")
-    torchtrt.save(trt_gm, trt_ep_path, inputs=args, kwargs_inputs=kwargs)
+    torchtrt.save(trt_gm, trt_ep_path, inputs=args, kwarg_inputs=kwargs)
     # Clean up model env
     torch._dynamo.reset()
 
@@ -201,7 +201,7 @@ def forward(self, x, b=5, c=None, d=None):
 
     # Save the module
     trt_ep_path = os.path.join(tempfile.gettempdir(), "compiled.ep")
-    torchtrt.save(trt_gm, trt_ep_path, inputs=args, kwargs_inputs=kwargs)
+    torchtrt.save(trt_gm, trt_ep_path, inputs=args, kwarg_inputs=kwargs)
     # Clean up model env
     torch._dynamo.reset()
 
@@ -288,7 +288,7 @@ def forward(self, x, b=None, c=None, d=None, e=[]):
     )
     # Save the module
     trt_ep_path = os.path.join(tempfile.gettempdir(), "compiled.ep")
-    torchtrt.save(trt_gm, trt_ep_path, inputs=args, kwargs_inputs=kwargs)
+    torchtrt.save(trt_gm, trt_ep_path, inputs=args, kwarg_inputs=kwargs)
     # Clean up model env
     torch._dynamo.reset()
 
@@ -375,7 +375,7 @@ def forward(self, x, b=None, c=None, d=None, e=[]):
     )
     # Save the module
     trt_ep_path = os.path.join(tempfile.gettempdir(), "compiled.ep")
-    torchtrt.save(trt_gm, trt_ep_path, inputs=args, kwargs_inputs=kwargs)
+    torchtrt.save(trt_gm, trt_ep_path, inputs=args, kwarg_inputs=kwargs)
     # Clean up model env
     torch._dynamo.reset()