Pass sample_kwargs properly during conversion.

- Reauthored models are getting more optional arguments: mask, lora, pixel_values.
- Without this change, sample_kwargs must be added with the exact order of arguments defined. For example, mask=None must be added before lora is added.
- With this change, sample_kwargs can omit optional args.
- Verified that lora conversion broken without this change is now working back with this change

PiperOrigin-RevId: 716414821

Author: ai-edge-bot

ai_edge_torch/_convert/conversion.py

@@ -109,21 +109,21 @@ def convert_signatures(
 
   _warn_training_modules(signatures)
 
-  def export(*args, **kwargs):
+  def export(**kwargs):
     nonlocal strict_export
     if strict_export == "auto":
       try:
-        exported_program = torch.export.export(*args, **kwargs, strict=True)
+        exported_program = torch.export.export(**kwargs, strict=True)
       except Exception:
         logging.warning(
             "torch.export.export(..., strict=True) failed. Retrying with"
             " strict=False"
         )
-        exported_program = torch.export.export(*args, **kwargs, strict=False)
+        exported_program = torch.export.export(**kwargs, strict=False)
     elif not strict_export:
-      exported_program = torch.export.export(*args, **kwargs, strict=False)
+      exported_program = torch.export.export(**kwargs, strict=False)
     else:
-      exported_program = torch.export.export(*args, **kwargs, strict=True)
+      exported_program = torch.export.export(**kwargs, strict=True)
 
     if hasattr(torch._decomp, "_decomp_table_to_post_autograd_aten"):
       # Available after torch 2.5.0: `_decomp_table_to_post_autograd_aten` is a
@@ -136,7 +136,12 @@ def export(*args, **kwargs):
     return exported_program
 
   exported_programs: torch.export.ExportedProgram = [
-      export(sig.module, sig.flat_args, dynamic_shapes=sig.dynamic_shapes)
+      export(
+          mod=sig.module,
+          args=sig.args,
+          kwargs=sig.kwargs,
+          dynamic_shapes=sig.dynamic_shapes,
+      )
       for sig in signatures
   ]
 

ai_edge_torch/_convert/signature.py

@@ -25,9 +25,9 @@
 class Signature:
   name: str
   module: torch.nn.Module
-  sample_args: tuple[torch.Tensor]
+  sample_args: tuple[torch.Tensor, ...]
   sample_kwargs: dict[str, torch.Tensor]
-  dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any]]] = None
+  dynamic_shapes: Optional[Union[Dict[str, Any], Tuple[Any, ...]]] = None
 
   @property
   def _normalized_sample_args_kwargs(self):
@@ -61,6 +61,16 @@ def flat_arg_names(self) -> list[str]:
     return names
 
   @property
-  def flat_args(self) -> tuple[Any]:
+  def flat_args(self) -> tuple[Any, ...]:
     args, kwargs = self._normalized_sample_args_kwargs
     return tuple([*args, *kwargs.values()])
+
+  @property
+  def args(self) -> tuple[Any, ...]:
+    args, _ = self._normalized_sample_args_kwargs
+    return args
+
+  @property
+  def kwargs(self) -> dict[str, Any]:
+    _, kwargs = self._normalized_sample_args_kwargs
+    return kwargs

ai_edge_torch/lowertools/common_utils.py

@@ -95,9 +95,9 @@ def _get_states(
     signatures: list[signature_module.Signature],
 ):
   for exported_program, signature in zip(exported_programs, signatures):
-    args, _ = exported_program.example_inputs
+    args, kwargs = exported_program.example_inputs
     # Calling this to get **all** the state including model buffers.
-    _flat_input_args = exported_program._graph_module_flat_inputs(args, {})
+    _flat_input_args = exported_program._graph_module_flat_inputs(args, kwargs)
     for tensor, input_spec in zip(
         _flat_input_args, exported_program.graph_signature.input_specs
     ):