supports integers for value_if_missing

Author: xadupre

_unittests/ut_investigate/test_input_observer.py

@@ -887,7 +887,7 @@ def forward(self, x=None, y=None):
         # self.assertEqual(2, len(args))
         # self.assertEqual(len([v for v in args.values() if v is not None]), 2)
 
-    def test_infer_dynamic_shapes_missing(self):
+    def test_infer_dynamic_shapes_missing_kwargs(self):
         class Model(torch.nn.Module):
             def forward(
                 self,
@@ -903,34 +903,34 @@ def forward(
 
         inputs = [
             dict(
-                input_ids=torch.ones((1, 282), dtype=torch.int64),
-                pixel_values=torch.ones((1, 3, 896, 896), dtype=torch.int64),
-                attention_mask=torch.ones((1, 282), dtype=torch.int64),
-                position_ids=torch.ones((1, 282), dtype=torch.int64),
-                token_type_ids=torch.ones((1, 282), dtype=torch.int64),
-                cache_position=torch.ones((282,), dtype=torch.int64),
+                input_ids=torch.ones((1, 28), dtype=torch.int64),
+                pixel_values=torch.ones((1, 3, 112, 112), dtype=torch.int64),
+                attention_mask=torch.ones((1, 28), dtype=torch.int64),
+                position_ids=torch.ones((1, 28), dtype=torch.int64),
+                token_type_ids=torch.ones((1, 28), dtype=torch.int64),
+                cache_position=torch.ones((28,), dtype=torch.int64),
             ),
             dict(
                 input_ids=torch.ones((1, 1), dtype=torch.int64),
-                attention_mask=torch.ones((1, 283), dtype=torch.int64),
+                attention_mask=torch.ones((1, 29), dtype=torch.int64),
                 position_ids=torch.ones((1, 1), dtype=torch.int64),
-                past_key_values=torch.rand((1, 1, 282, 32)),
+                past_key_values=torch.rand((1, 1, 28, 32)),
                 token_type_ids=torch.ones((1, 1), dtype=torch.int64),
                 cache_position=torch.ones((1,), dtype=torch.int64),
             ),
             dict(
                 input_ids=torch.ones((1, 1), dtype=torch.int64),
-                attention_mask=torch.ones((1, 284), dtype=torch.int64),
+                attention_mask=torch.ones((1, 30), dtype=torch.int64),
                 position_ids=torch.ones((1, 1), dtype=torch.int64),
-                past_key_values=torch.rand((1, 1, 283, 32)),
+                past_key_values=torch.rand((1, 1, 29, 32)),
                 token_type_ids=torch.ones((1, 1), dtype=torch.int64),
                 cache_position=torch.ones((1,), dtype=torch.int64),
             ),
         ]
 
         model = Model()
         observer = InputObserver(
-            value_if_missing=dict(pixel_values=torch.empty((0, 3, 896, 896)))
+            value_if_missing=dict(pixel_values=torch.empty((0, 3, 112, 112)))
         )
         with observer(model):
             for kwargs in inputs:
@@ -948,6 +948,132 @@ def forward(
             "cache_position": {0: cst},
         }
         self.assertEqual(expected, shapes)
+        kwargs = observer.infer_arguments()
+        self.assertEqual(list(expected), list(kwargs))
+        self.assertEqual((0, 3, 112, 112), kwargs["pixel_values"].shape)
+
+    def test_infer_dynamic_shapes_missing_args(self):
+        class Model(torch.nn.Module):
+            def forward(
+                self,
+                input_ids=None,
+                pixel_values=None,
+                attention_mask=None,
+                past_key_values=None,
+            ):
+                return input_ids
+
+        inputs = [
+            (
+                torch.ones((1, 28), dtype=torch.int64),
+                torch.ones((1, 3, 112, 112), dtype=torch.int64),
+                torch.ones((1, 28), dtype=torch.int64),
+            ),
+            (
+                torch.ones((1, 1), dtype=torch.int64),
+                None,
+                torch.ones((1, 29), dtype=torch.int64),
+                torch.rand((1, 1, 28, 32)),
+            ),
+            (
+                torch.ones((1, 1), dtype=torch.int64),
+                None,
+                torch.ones((1, 30), dtype=torch.int64),
+                torch.rand((1, 1, 29, 32)),
+            ),
+        ]
+
+        model = Model()
+        observer = InputObserver(
+            value_if_missing={1: torch.empty((0, 3, 112, 112), dtype=torch.int64)}
+        )
+        with observer(model):
+            for args in inputs:
+                model(*args)
+
+        shapes = observer.infer_dynamic_shapes(set_batch_dimension_for=True)
+        cst = torch.export.Dim.DYNAMIC
+        expected = ({0: cst, 1: cst}, {0: cst}, {0: cst, 1: cst}, {0: cst, 2: cst})
+        self.assertEqual(expected, shapes)
+        args = observer.infer_arguments()
+        self.assertEqual(len(expected), len(args))
+        self.assertEqual((0, 3, 112, 112), args[1].shape)
+
+    def test_infer_dynamic_shapes_missing_kwargs_nested(self):
+        class Model(torch.nn.Module):
+            def forward(
+                self,
+                input_ids=None,
+                pixel_values=None,
+                attention_mask=None,
+                position_ids=None,
+                past_key_values=None,
+                token_type_ids=None,
+                cache_position=None,
+            ):
+                return input_ids
+
+        inputs = [
+            dict(
+                input_ids=torch.ones((1, 28), dtype=torch.int64),
+                pixel_values=(
+                    torch.ones((1, 3, 112, 112), dtype=torch.int64),
+                    torch.ones((1, 3, 112, 112), dtype=torch.int64),
+                ),
+                attention_mask=torch.ones((1, 28), dtype=torch.int64),
+                position_ids=torch.ones((1, 28), dtype=torch.int64),
+                token_type_ids=torch.ones((1, 28), dtype=torch.int64),
+                cache_position=torch.ones((28,), dtype=torch.int64),
+            ),
+            dict(
+                input_ids=torch.ones((1, 1), dtype=torch.int64),
+                attention_mask=torch.ones((1, 29), dtype=torch.int64),
+                position_ids=torch.ones((1, 1), dtype=torch.int64),
+                past_key_values=torch.rand((1, 1, 28, 32)),
+                token_type_ids=torch.ones((1, 1), dtype=torch.int64),
+                cache_position=torch.ones((1,), dtype=torch.int64),
+            ),
+            dict(
+                input_ids=torch.ones((1, 1), dtype=torch.int64),
+                attention_mask=torch.ones((1, 30), dtype=torch.int64),
+                position_ids=torch.ones((1, 1), dtype=torch.int64),
+                past_key_values=torch.rand((1, 1, 29, 32)),
+                token_type_ids=torch.ones((1, 1), dtype=torch.int64),
+                cache_position=torch.ones((1,), dtype=torch.int64),
+            ),
+        ]
+
+        model = Model()
+        observer = InputObserver(
+            value_if_missing=dict(
+                pixel_values=(
+                    torch.empty((0, 3, 112, 112), dtype=torch.int64),
+                    torch.empty((0, 3, 112, 112), dtype=torch.int64),
+                )
+            )
+        )
+        with observer(model):
+            for kwargs in inputs:
+                model(**kwargs)
+
+        shapes = observer.infer_dynamic_shapes(set_batch_dimension_for=True)
+        cst = torch.export.Dim.DYNAMIC
+        expected = {
+            "input_ids": {0: cst, 1: cst},
+            "pixel_values": ({0: cst}, {0: cst}),
+            "attention_mask": {0: cst, 1: cst},
+            "position_ids": {0: cst, 1: cst},
+            "past_key_values": {0: cst, 2: cst},
+            "token_type_ids": {0: cst, 1: cst},
+            "cache_position": {0: cst},
+        }
+        self.assertEqual(expected, shapes)
+        kwargs = observer.infer_arguments()
+        self.assertEqual(list(expected), list(kwargs))
+        self.assertIsInstance(kwargs["pixel_values"], tuple)
+        self.assertEqual(2, len(kwargs["pixel_values"]))
+        self.assertEqual((0, 3, 112, 112), kwargs["pixel_values"][0].shape)
+        self.assertEqual((0, 3, 112, 112), kwargs["pixel_values"][1].shape)
 
     def test_io_captured_kwargs_kwargs(self):
         class Model(torch.nn.Module):

_unittests/ut_investigate/test_input_observer_transformers.py

@@ -280,7 +280,9 @@ def forward(
 
         model = Model()
         observer = InputObserver(
-            value_if_missing=dict(pixel_values=torch.empty((0, 3, 896, 896)))
+            value_if_missing=dict(
+                pixel_values=torch.empty((0, 3, 896, 896), dtype=torch.int64)
+            )
         )
         with (
             register_additional_serialization_functions(patch_transformers=True),

onnx_diagnostic/investigate/input_observer.py

@@ -87,7 +87,7 @@ def _infer_dynamic_dimensions(
     unique_ranks = {len(shape) for shape in shape_list}
     torch._check(
         len(unique_ranks) == 1,
-        lambda: "all shapes in shape_list must have the same rank",
+        lambda: f"All shapes in shape_list must have the same rank but {shape_list=}.",
     )
     rank = unique_ranks.pop()
     dynamic = []
@@ -129,7 +129,6 @@ def __init__(
         self.args = args
         self.kwargs = kwargs
         self.flat_list, self.spec = torch.utils._pytree.tree_flatten((args, kwargs))
-        self.n_tensors = sum(t is not None for t in self.flat_list)
         self._position_to_args_kwargs: list[int | str] | None = None
         self._n_tensors_for_args_kwargs: dict[int | str, int] | None = None
         self.cst_kwargs = cst_kwargs.copy()
@@ -290,7 +289,7 @@ class InputObserverInfo:
             to be the same in the ordered dictionaries `add_inputs` receive.
         default_values: Default values defined by the signature of the function,
             any value equal to that is ignored to simplify the export.
-        value_if_missing: If a named argument (in kwargs) is missing,
+        value_if_missing: If an argument is missing,
             a default value will be taken in this dictionary,
             this is used when after the prefill step, an argument
             disappears (such as `pixel_values`) and another one
@@ -308,7 +307,7 @@ def __init__(
         self,
         signature_names: list[str],
         default_values: dict[str, int | bool | str | float],
-        value_if_missing: dict[str, Any],
+        value_if_missing: dict[str | int, Any],
         args_name_and_position: tuple[str, int] | None,
         kwargs_name: str | None,
     ):
@@ -351,21 +350,47 @@ def add_inputs(self, args: tuple[Any, ...], kwargs: dict[str, Any]):
 
         # adds value_if_missing attributes
         for k, v in self.value_if_missing.items():
-            if k not in kwargs:
-                # Validate that `value_if_missing` keys are compatible
-                # with the observed signature.
-                # If the function does not accept **kwargs,
-                # all value_if_missing keys must be
-                # present in the observed signature names.
-                if k not in self.signature_names and not self.kwargs_name:
+            if isinstance(k, str):
+                if k not in kwargs:
+                    # Validate that `value_if_missing` keys are compatible
+                    # with the observed signature.
+                    # If the function does not accept **kwargs,
+                    # all value_if_missing keys must be
+                    # present in the observed signature names.
+                    if k not in self.signature_names and not self.kwargs_name:
+                        raise ValueError(
+                            f"Unexpected keyword argument {k!r} "
+                            f"provided as a value_if_missing input "
+                            "for a function that does not accept it. "
+                            f"All value_if_missing keys must "
+                            f"be in the observed signature: {tuple(self.signature_names)}."
+                        )
+                    kwargs[k] = v
+            elif isinstance(k, int):
+                if k >= len(self.signature_names):
                     raise ValueError(
-                        f"Unexpected keyword argument '{k}' "
+                        f"Unexpected keyword argument {k=} "
                         f"provided as a value_if_missing input "
                         "for a function that does not accept it. "
-                        f"All value_if_missing keys must "
+                        f"All value_if_missing indices must "
                         f"be in the observed signature: {tuple(self.signature_names)}."
                     )
-                kwargs[k] = v
+                if k >= len(args):
+                    raise NotImplementedError(
+                        f"Unexpected keyword argument {k=} "
+                        f"provided as a value_if_missing input "
+                        "for a function that does not accept it. "
+                        f"All value_if_missing indices must "
+                        f"be in the observed signature: {tuple(self.signature_names)}, "
+                        f"only {len(args)} were given."
+                    )
+                list_args = list(args)
+                list_args[k] = v
+                args = tuple(list_args)
+            else:
+                raise TypeError(
+                    f"Unexepcted type {type(k)} for a missing value. The key is {k!r}."
+                )
 
         # kwargs may come in a different order each time.
         # dictionaries are ordered and torch.export.export expects
@@ -793,7 +818,7 @@ class InputObserver:
     export arguments.
 
     Args:
-        value_if_missing: If a named argument (in kwargs) is missing,
+        value_if_missing: If an argument is missing,
             a default value will be taken in this dictionary,
             this is used when after the prefill step, an argument
             disappears (such as `pixel_values`) and another one
@@ -877,7 +902,7 @@ class InputObserver:
     :ref:`l-plot-gemma3-tiny-export-input-observer`.
     """
 
-    def __init__(self, value_if_missing: dict[str, Any] | None = None):
+    def __init__(self, value_if_missing: dict[str | int, Any] | None = None):
         self.info: InputObserverInfo | None = None  # type: ignore[annotation-unchecked]
         self.value_if_missing = value_if_missing or {}