json ambiguities

Author: xadupre

_doc/recipes/plot_dynamic_shapes_json.py

@@ -0,0 +1,104 @@
+"""
+JSON returns list when the original dynamic shapes are list or tuple
+====================================================================
+
+Dynamic Shapes After JSON
++++++++++++++++++++++++++
+"""
+
+import json
+import pprint
+import torch
+from onnx_diagnostic import doc
+from onnx_diagnostic.helpers import string_type
+from onnx_diagnostic.helpers.cache_helper import make_dynamic_cache
+from onnx_diagnostic.export.shape_helper import all_dynamic_shape_from_inputs
+
+bsize, nheads, slen, dim = 2, 1, 30, 96
+
+inputs = dict(
+    input_mask_position=(
+        torch.randint(15, size=(2, 3), dtype=torch.int64),
+        torch.randint(1, size=(2, 33), dtype=torch.int64),
+        torch.arange(3, dtype=torch.int64),
+    ),
+    past_key_values=make_dynamic_cache(
+        [(torch.randn(bsize, nheads, slen, dim), torch.randn(bsize, nheads, slen, dim))]
+    ),
+)
+
+print(string_type(inputs, with_shape=True))
+
+# %%
+# Function :func:`onnx_diagnostic.export.shape_helper.all_dynamic_shape_from_inputs`
+# produces the corresponding dynamic shapes assuming they are all dynamic.
+ds = all_dynamic_shape_from_inputs(inputs)
+pprint.pprint(ds)
+
+# %%
+# Converted into JSON.
+
+json_str = json.dumps(ds, indent=2, ensure_ascii=False)
+print(json_str)
+
+# %%
+# Restoration.
+ds2 = json.loads(json_str)
+pprint.pprint(ds2)
+
+# %%
+# tuple are replaced by list.
+
+# The trick
+# +++++++++
+
+
+def flatten_unflatten_like_dynamic_shapes(obj):
+    if isinstance(obj, torch.Tensor):
+        return obj
+    flat, spec = torch.utils._pytree.tree_flatten(obj)
+    start = 0
+    end = 0
+    subtrees = []
+    for subspec in spec.children_specs:
+        end += subspec.num_leaves
+        value = subspec.unflatten(flat[start:end])
+        value = flatten_unflatten_like_dynamic_shapes(value)
+        subtrees.append(value)
+        start = end
+    if spec.type is dict or spec.context:
+        return dict(zip(spec.context, subtrees))
+    if spec.type is tuple:
+        return tuple(subtrees)
+    return subtrees
+
+
+def _align(inputs, ds):
+    if isinstance(inputs, torch.Tensor):
+        return ds
+    if isinstance(inputs, tuple):
+        return tuple(_align(o, d) for o, d in zip(inputs, ds))
+    if isinstance(inputs, list):
+        return [_align(o, d) for o, d in zip(inputs, ds)]
+    if isinstance(inputs, dict):
+        return {k: _align(inputs[k], d) for k, d in ds.items()}
+    raise TypeError(f"Unexpected types inputs is {type(inputs)}, ds is {type(ds)}")
+
+
+def fix_dynamic_shapes(inputs, dynamic_shapes):
+    flat_unflat_inputs = flatten_unflatten_like_dynamic_shapes(inputs)
+    return _align(flat_unflat_inputs, dynamic_shapes)
+
+
+fixed_ds = fix_dynamic_shapes(inputs, ds2)
+pprint.pprint(fixed_ds)
+
+# %%
+# The code changed tuple into list as expected.
+assert isinstance(ds2["input_mask_position"], list)
+assert isinstance(fixed_ds["input_mask_position"], tuple)
+
+
+# %%
+
+doc.plot_legend("dynamic shapes\nto json\nfrom json", "torch.export.export", "green")

onnx_diagnostic/helpers/cache_helper.py

@@ -19,7 +19,7 @@ def flatten_unflatten_for_dynamic_shapes(
         :func:`torch.export.export` only considers the values,
         the context gives the dictionary keys but it is not expressed
         in the dynamic shapes, these specifications seems to be different
-        for the strict and non strict mode.
+        for the strict and non strict mode. It also preserves tuple.
     :param change_function: to modifies the tensor in the structure itself,
         like replace them by a shape
     :return: the serialized object
@@ -38,9 +38,12 @@ def flatten_unflatten_for_dynamic_shapes(
         )
         subtrees.append(value)
         start = end
-    if use_dict and (spec.type is dict or spec.context):
-        # This a dictionary.
-        return dict(zip(spec.context, subtrees))
+    if use_dict:
+        if spec.type is dict or spec.context:
+            # This a dictionary.
+            return dict(zip(spec.context, subtrees))
+        if spec.type is tuple:
+            return tuple(subtrees)
     # This is a list.
     return subtrees