Using string instead of Dim

Author: xadupre

onnx_diagnostic/tasks/image_text_to_text.py

@@ -52,19 +52,19 @@ def get_inputs(
     :return: dictionary
     """
     batch = torch.export.Dim("batch", min=1, max=1024)
-    seq_length = torch.export.Dim("seq_length", min=1, max=4096)
-    cache_length = torch.export.Dim("cache_length", min=1, max=4096)
-    images = torch.export.Dim("images", min=1, max=4096)
+    seq_length = "seq_length"  # torch.export.Dim("seq_length", min=1, max=4096)
+    cache_length = "cache_length"  # torch.export.Dim("cache_length", min=1, max=4096)
+    images = "images"  # torch.export.Dim("images", min=1, max=4096)
 
     shapes = {
         "input_ids": {0: batch, 1: seq_length},
         "attention_mask": {
             0: batch,
-            1: torch.export.Dim.DYNAMIC,  # cache_length + seq_length
+            1: "cache+seq",  # cache_length + seq_length
         },
         "position_ids": {
             0: batch,
-            1: torch.export.Dim.DYNAMIC,  # cache_length + seq_length
+            1: "cache+seq",  # cache_length + seq_length
         },
         "past_key_values": [
             [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],

onnx_diagnostic/tasks/text2text_generation.py

@@ -59,14 +59,14 @@ def get_inputs(
         encoder_outputs:dict(last_hidden_state:T1s1x16x512)
     """
     batch = torch.export.Dim("batch", min=1, max=1024)
-    seq_length = torch.export.Dim("seq_length", min=1, max=4096)
-    cache_length = torch.export.Dim("cache_length", min=1, max=4096)
-    cache_length2 = torch.export.Dim("cache_length2", min=1, max=4096)
+    seq_length = "seq_length"  # torch.export.Dim("seq_length", min=1, max=4096)
+    cache_length = "cache_length_key"  # torch.export.Dim("cache_length", min=1, max=4096)
+    cache_length2 = "cache_length_val"  # torch.export.Dim("cache_length2", min=1, max=4096)
 
     shapes = {
         "input_ids": {0: batch, 1: seq_length},
-        "decoder_input_ids": {0: batch, 1: torch.export.Dim.DYNAMIC},
-        "attention_mask": {0: batch, 1: torch.export.Dim.DYNAMIC},
+        "decoder_input_ids": {0: batch, 1: "seq_ids"},
+        "attention_mask": {0: batch, 1: "seq_mask"},
         # "cache_position": {0: batch, 1: torch.export.Dim.DYNAMIC},
         "past_key_values": [
             [

onnx_diagnostic/tasks/text_classification.py

@@ -35,7 +35,7 @@ def get_inputs(
         attention_mask:T7s1x13[1,1:A1.0])
     """
     batch = torch.export.Dim("batch", min=1, max=1024)
-    seq_length = torch.export.Dim("sequence_length", min=1, max=1024)
+    seq_length = "seq_length"  # torch.export.Dim("sequence_length", min=1, max=1024)
     shapes = {
         "input_ids": {0: batch, 1: seq_length},
         "token_type_ids": {0: batch, 1: seq_length},

onnx_diagnostic/tasks/text_generation.py

@@ -84,8 +84,8 @@ def get_inputs(
     :return: dictionary
     """
     batch = torch.export.Dim("batch", min=1, max=1024)
-    seq_length = torch.export.Dim("seq_length", min=1, max=4096)
-    cache_length = torch.export.Dim("cache_length", min=1, max=4096)
+    seq_length = "seq_length"  # torch.export.Dim("seq_length", min=1, max=4096)
+    cache_length = "cache_length"  # torch.export.Dim("cache_length", min=1, max=4096)
 
     if config is not None and config.__class__.__name__ == "FalconMambaConfig":
         seq_length_multiple = 8
@@ -101,11 +101,11 @@ def get_inputs(
             "input_ids": {0: batch, 1: torch.export.Dim.DYNAMIC},
             "attention_mask": {
                 0: batch,
-                1: torch.export.Dim.DYNAMIC,  # cache_length + seq_length
+                1: "cache+seq",  # cache_length + seq_length
             },
             "cache_position": {
                 0: batch,
-                1: torch.export.Dim.DYNAMIC,  # cache_length + seq_length
+                1: "cache+seq",  # cache_length + seq_length
             },
             "cache_params": [
                 [{0: batch} for _ in range(num_hidden_layers)],
@@ -145,11 +145,11 @@ def get_inputs(
         "input_ids": {0: batch, 1: seq_length},
         "attention_mask": {
             0: batch,
-            1: torch.export.Dim.DYNAMIC,  # cache_length + seq_length
+            1: "cache+seq",  # cache_length + seq_length
         },
         "position_ids": {
             0: batch,
-            1: torch.export.Dim.DYNAMIC,  # cache_length + seq_length
+            1: "cache+seq",  # cache_length + seq_length
         },
         "past_key_values": [
             [{0: batch, 2: cache_length} for _ in range(num_hidden_layers)],