Add desired_values to change_dynamic_dimensions (#45)

* Add desired_values to change_dynamic_dimensions

* change

Author: xadupre

CHANGELOGS.rst

@@ -1,6 +1,11 @@
 Change Logs
 ===========
 
+0.4.0
++++++
+
+* :pr:`45`: improve change_dynamic_dimension to fix some dimensions
+
 0.3.0
 +++++
 

_unittests/ut_export/test_dynamic_shapes.py

@@ -691,6 +691,18 @@ def test_couple_input_ds_change_dynamic_dimensions(self):
         self.assertEqual((3, 5, 8), new_input["A"].shape)
         self.assertEqual((3, 10), new_input["B"].shape)
 
+    def test_couple_input_ds_change_dynamic_dimensions_fixed(self):
+        T257 = torch.arange(2 * 5 * 7).reshape((2, 5, 7))
+        T29 = torch.arange(2 * 9).reshape((2, 9))
+        inst = CoupleInputsDynamicShapes(
+            (),
+            {"A": T257, "B": T29},
+            {"A": {0: "batch", 2: "last"}, "B": {0: "batch", 1: "seq"}},
+        )
+        new_input = inst.change_dynamic_dimensions({"seq": 50, "batch": 1})
+        self.assertEqual((1, 5, 8), new_input["A"].shape)
+        self.assertEqual((1, 50), new_input["B"].shape)
+
 
 if __name__ == "__main__":
     unittest.main(verbosity=2)

onnx_diagnostic/export/dynamic_shapes.py

@@ -243,8 +243,8 @@ def _generic_walker_step(cls, processor: Callable, inputs, ds):
         return cls._generic_walker_step(processor, flat, ds)
 
     class ChangeDimensionProcessor:
-        def __init__(self):
-            self.mapping = {}
+        def __init__(self, desired_values):
+            self.mapping = desired_values or {}
 
         def _build_new_shape(
             self, shape: Tuple[int, ...], ds: Dict[int, Any]
@@ -310,7 +310,7 @@ def __call__(self, inputs, ds):
             new_shape = self._build_new_shape(inputs.shape, ds)
             return self._build_new_tensor(inputs, new_shape)
 
-    def change_dynamic_dimensions(self):
+    def change_dynamic_dimensions(self, desired_values: Optional[Dict[str, int]] = None):
         """
         A model exported with dynamic shapes is not necessarily dynamic
         just because the user specified dynamic shapes. The algorithm
@@ -320,6 +320,9 @@ def change_dynamic_dimensions(self):
         for the dimension than the first ones, assuming they were used to export
         the model.
 
+        :param desired_values: to fixed named dimension to have the desired value
+        :return: new inputs
+
         Example:
 
         .. runpython::
@@ -340,7 +343,7 @@ def change_dynamic_dimensions(self):
             print("before:", string_type(kwargs, with_shape=True))
             print("-after:", string_type(new_kwargs, with_shape=True))
         """
-        return self._generic_walker(self.ChangeDimensionProcessor())
+        return self._generic_walker(self.ChangeDimensionProcessor(desired_values))
 
 
 class ModelInputs: