onnx-diagnostic 0.8.1__py3-none-any.whl → 0.8.3__py3-none-any.whl

onnx_diagnostic/helpers/dot_helper.py

@@ -0,0 +1,210 @@
+from typing import Dict, Set
+import onnx
+import onnx.numpy_helper as onh
+from .onnx_helper import onnx_dtype_name, pretty_onnx
+
+
+def _get_hidden_inputs(graph: onnx.GraphProto) -> Set[str]:
+    hidden = set()
+    memo = (
+        {i.name for i in graph.initializer}
+        | {i.values.name for i in graph.sparse_initializer}
+        | {i.name for i in graph.input}
+    )
+    for node in graph.node:
+        for i in node.input:
+            if i not in memo:
+                hidden.add(i)
+        for att in node.attribute:
+            if att.type == onnx.AttributeProto.GRAPH and att.g:
+                hid = _get_hidden_inputs(att.g)
+                less = set(h for h in hid if h not in memo)
+                hidden |= less
+        memo |= set(node.output)
+    return hidden
+
+
+def _make_node_label(node: onnx.NodeProto, tiny_inits: Dict[str, str]) -> str:
+    els = [f"{node.domain}.\\n{node.op_type}" if node.domain else node.op_type, "("]
+    ee = [tiny_inits.get(i, ".") if i else "" for i in node.input]
+    for att in node.attribute:
+        if att.name == "to":
+            ee.append(f"{att.name}={onnx_dtype_name(att.i)}")
+        elif att.name in {"to", "axis", "value_int", "stash_type", "start", "end"}:
+            ee.append(f"{att.name}={att.i}")
+        elif att.name in {"value_float"}:
+            ee.append(f"{att.name}={att.f}")
+        elif att.name in {"value_floats"}:
+            ee.append(f"{att.name}={att.floats}")
+        elif att.name in {"value_ints", "perm"}:
+            ee.append(f"{att.name}={att.ints}")
+    els.append(", ".join(ee))
+    els.append(")")
+    if node.op_type == "Constant":
+        els.extend([" -> ", node.output[0]])
+    return "".join(els)
+
+
+def _make_edge_label(value_info: onnx.ValueInfoProto, multi_line: bool = False) -> str:
+    itype = value_info.type.tensor_type.elem_type
+    if itype == onnx.TensorProto.UNDEFINED:
+        return ""
+    shape = tuple(
+        d.dim_param if d.dim_param else d.dim_value
+        for d in value_info.type.tensor_type.shape.dim
+    )
+    res = [
+        str(a)
+        for a in [("?" if isinstance(s, str) and s.startswith("unk") else s) for s in shape]
+    ]
+    sshape = ",".join(res)
+    if multi_line and len(sshape) > 30:
+        sshape = ",\\n".join(res)
+    return f"{onnx_dtype_name(itype)}({sshape})"
+
+
+def to_dot(model: onnx.ModelProto) -> str:
+    """
+    Converts a model into a dot graph.
+    Here is an example:
+
+    .. gdot::
+        :script: DOT-SECTION
+        :process:
+
+        from onnx_diagnostic.helpers.dot_helper import to_dot
+        from onnx_diagnostic.export.api import to_onnx
+        from onnx_diagnostic.torch_export_patches import torch_export_patches
+        from onnx_diagnostic.torch_models.hghub import get_untrained_model_with_inputs
+
+        data = get_untrained_model_with_inputs("arnir0/Tiny-LLM")
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        with torch_export_patches(patch_transformers=True):
+            em = to_onnx(model, inputs, dynamic_shapes=ds, exporter="custom")
+        dot = to_dot(em.model_proto)
+        print("DOT-SECTION", dot)
+
+    Or this one obtained with :func:`torch.onnx.export`.
+
+    .. gdot::
+        :script: DOT-SECTION
+        :process:
+
+        from onnx_diagnostic.helpers.dot_helper import to_dot
+        from onnx_diagnostic.export.api import to_onnx
+        from onnx_diagnostic.torch_export_patches import torch_export_patches
+        from onnx_diagnostic.torch_models.hghub import get_untrained_model_with_inputs
+
+        data = get_untrained_model_with_inputs("arnir0/Tiny-LLM")
+        model, inputs, ds = data["model"], data["inputs"], data["dynamic_shapes"]
+        with torch_export_patches(patch_transformers=True):
+            em = to_onnx(model, kwargs=inputs, dynamic_shapes=ds, exporter="onnx-dynamo")
+        dot = to_dot(em.model_proto)
+        print("DOT-SECTION", dot)
+    """
+    _unique: Dict[int, int] = {}
+
+    def _mkn(obj: object) -> int:
+        id_obj = id(obj)
+        if id_obj in _unique:
+            return _unique[id_obj]
+        i = len(_unique)
+        _unique[id_obj] = i
+        return i
+
+    model = onnx.shape_inference.infer_shapes(model)
+
+    op_type_colors = {
+        "Shape": "#d2a81f",
+        "MatMul": "#ee9999",
+        "Transpose": "#ee99ee",
+        "Reshape": "#eeeeee",
+        "Squeeze": "#eeeeee",
+        "Unsqueeze": "#eeeeee",
+    }
+
+    edge_label = {}
+    for val in model.graph.value_info:
+        edge_label[val.name] = _make_edge_label(val, multi_line=True)
+
+    rows = [
+        "digraph {",
+        (
+            "  graph [rankdir=TB, splines=true, overlap=false, nodesep=0.2, "
+            "ranksep=0.2, fontsize=8];"
+        ),
+        '  node [style="rounded,filled", color="#888888", fontcolor="#222222", shape=box];',
+        "  edge [arrowhead=vee, fontsize=7, labeldistance=-5, labelangle=0];",
+    ]
+    inputs = list(model.graph.input)
+    outputs = list(model.graph.output)
+    nodes = list(model.graph.node)
+    inits = list(model.graph.initializer)
+    tiny_inits = {}
+    name_to_ids = {}
+    for inp in inputs:
+        if not inp.name:
+            continue
+        lab = _make_edge_label(inp)
+        rows.append(f'  I_{_mkn(inp)} [label="{inp.name}\\n{lab}", fillcolor="#aaeeaa"];')
+        name_to_ids[inp.name] = f"I_{_mkn(inp)}"
+        edge_label[inp.name] = _make_edge_label(inp, multi_line=True)
+    for init in inits:
+        shape = tuple(init.dims)
+        if len(shape) == 0 or (len(shape) == 1 and shape[0] < 10):
+            a = onh.to_array(init)
+            tiny_inits[init.name] = (
+                str(a) if len(shape) == 0 else f"[{', '.join([str(i) for i in a])}]"
+            )
+        else:
+            ls = f"{onnx_dtype_name(init.data_type)}({', '.join(map(str,shape))})"
+            rows.append(f'  i_{_mkn(init)} [label="{init.name}\\n{ls}", fillcolor="#cccc00"];')
+            name_to_ids[init.name] = f"i_{_mkn(init)}"
+            edge_label[init.name] = ls
+    for node in nodes:
+        color = op_type_colors.get(node.op_type, "#cccccc")
+        label = _make_node_label(node, tiny_inits)
+        rows.append(f'  {node.op_type}_{_mkn(node)} [label="{label}", fillcolor="{color}"];')
+        name_to_ids.update({o: f"{node.op_type}_{_mkn(node)}" for o in node.output if o})
+
+    # nodes
+    done = set()
+    for node in nodes:
+        names = list(node.input)
+        for i in names:
+            if not i or i in tiny_inits:
+                continue
+            if i not in name_to_ids:
+                raise ValueError(f"Unable to find {i!r}\n{pretty_onnx(model)}")
+            edge = name_to_ids[i], f"{node.op_type}_{_mkn(node)}"
+            if edge in done:
+                continue
+            done.add(edge)
+            lab = edge_label.get(i, "")
+            if lab:
+                ls = ",".join([f'label="{lab}"'])
+                lab = f" [{ls}]"
+            rows.append(f"  {edge[0]} -> {edge[1]}{lab};")
+        if node.op_type in {"Scan", "Loop", "If"}:
+            unique = set()
+            for att in node.attribute:
+                if att.type == onnx.AttributeProto.GRAPH:
+                    unique |= _get_hidden_inputs(att.g)
+            for i in unique:
+                edge = name_to_ids[i], _mkn(node)  # type: ignore[assignment]
+                if edge in done:
+                    continue
+                done.add(edge)
+                rows.append(f"  {edge[0]} -> {edge[1]} [style=dotted];")
+
+    # outputs
+    for out in outputs:
+        if not out.name:
+            continue
+        lab = _make_edge_label(out)
+        rows.append(f'  O_{_mkn(out)} [label="{out.name}\\n{lab}", fillcolor="#aaaaee"];')
+        edge = name_to_ids[out.name], f"O_{_mkn(out)}"
+        rows.append(f"  {edge[0]} -> {edge[1]};")
+
+    rows.append("}")
+    return "\n".join(rows)

onnx_diagnostic/helpers/helper.py

@@ -529,16 +529,20 @@ def string_type(
                 return "OV(NO-NUMPY:FIXIT)"
             if verbose:
                 print(f"[string_type] V4:{type(obj)}")
-            return f"OV({string_type(t, with_shape=with_shape, with_min_max=with_min_max)})"
+            dev = ("G" if obj.device_name() == "Cuda" else "C") if with_device else ""
+            return (
+                f"{dev}OV({string_type(t, with_shape=with_shape, with_min_max=with_min_max)})"
+            )
         dt = obj.element_type()
         shape = obj.shape()
+        dev = ("G" if obj.device_name() == "Cuda" else "C") if with_device else ""
         if with_shape:
             if verbose:
                 print(f"[string_type] V5:{type(obj)}")
-            return f"OV{dt}s{'x'.join(map(str, shape))}"
+            return f"{dev}OV{dt}s{'x'.join(map(str, shape))}"
         if verbose:
             print(f"[string_type] V6:{type(obj)}")
-        return f"OV{dt}r{len(shape)}"
+        return f"{dev}OV{dt}r{len(shape)}"
 
     # others classes
 
@@ -787,6 +791,8 @@ def string_type(
         return f"ultralytics.{obj.__class__.__name__}(...)"
     if obj.__class__.__name__ == "FakeTensorMode":
         return f"{obj}"
+    if obj.__class__.__name__ == "FakeTensorContext":
+        return "FakeTensorContext(...)"
 
     if verbose:
         print(f"[string_type] END:{type(obj)}")
@@ -988,7 +994,7 @@ def max_diff(
     _index: int = 0,
     allow_unique_tensor_with_list_of_one_element: bool = True,
     hist: Optional[Union[bool, List[float]]] = None,
-) -> Dict[str, Union[float, int, Tuple[int, ...]]]:
+) -> Dict[str, Union[float, int, Tuple[Any, ...]]]:
     """
     Returns the maximum discrepancy.
 
@@ -1013,6 +1019,7 @@ def max_diff(
         output, this number will be the number of elements
         of this output
     * dnan: difference in the number of nan
+    * dev: tensor on the same device, if applicable
 
     You may use :func:`string_diff` to display the discrepancies in one string.
     """
@@ -1165,7 +1172,7 @@ def max_diff(
 
         if verbose >= 6:
             print(f"[max_diff] list,tuple,6: {string_type(expected)} ? {string_type(got)}")
-        am, rm, sm, n, dn, drep = 0, 0, 0.0, 0.0, 0, None
+        am, rm, sm, n, dn, drep, dd = 0, 0, 0.0, 0.0, 0, None, None
         for ip, (e, g) in enumerate(zip(expected, got)):
             d = max_diff(
                 e,
@@ -1197,7 +1204,15 @@ def max_diff(
                 else:
                     for k, v in d["rep"].items():
                         drep[k] += v
+            if "dev" in d and d["dev"] is not None:
+                if dd is None:
+                    dd = d["dev"]
+                else:
+                    dd += d["dev"]  # type: ignore[operator]
+
         res = dict(abs=am, rel=rm, sum=sm, n=n, dnan=dn)
+        if dd is not None:
+            res["dev"] = dd
         if drep:
             res["rep"] = drep
         return res  # type: ignore
@@ -1231,33 +1246,42 @@ def max_diff(
     import torch
 
     if isinstance(expected, np.ndarray) or isinstance(got, np.ndarray):
+        dev = None
         if isinstance(expected, torch.Tensor):
             from .torch_helper import to_numpy
 
+            dev = 0 if expected.device.type == "cpu" else 1
             expected = to_numpy(expected)
         if isinstance(got, torch.Tensor):
             from .torch_helper import to_numpy
 
+            dev = 0 if got.device.type == "cpu" else 1
             got = to_numpy(got)
         if verbose >= 6:
             print(f"[max_diff] tensor: {string_type(expected)} ? {string_type(got)}")
 
         if _index < begin or (end != -1 and _index >= end):
             # out of boundary
-            return dict(abs=0.0, rel=0.0, sum=0.0, n=0.0, dnan=0)
+            res = dict(abs=0.0, rel=0.0, sum=0.0, n=0.0, dnan=0)
+            if dev is not None:
+                res["dev"] = dev  # type: ignore[operator]
+            return res  # type: ignore[return-value]
         if isinstance(expected, (int, float)):
             if isinstance(got, np.ndarray) and len(got.shape) == 0:
                 got = float(got)
             if isinstance(got, (int, float)):
                 if expected == got:
                     return dict(abs=0.0, rel=0.0, sum=0.0, n=0.0, dnan=0)
-                return dict(
+                res = dict(
                     abs=abs(expected - got),
                     rel=abs(expected - got) / (abs(expected) + 1e-5),
                     sum=abs(expected - got),
                     n=1,
                     dnan=0,
                 )
+                if dev is not None:
+                    res["dev"] = dev
+                return res  # type: ignore[return-value]
             return dict(abs=np.inf, rel=np.inf, sum=np.inf, n=np.inf, dnan=np.inf)
         if expected.dtype in (np.complex64, np.complex128):
             if got.dtype == expected.dtype:
@@ -1337,6 +1361,8 @@ def max_diff(
         res: Dict[str, float] = dict(  # type: ignore
             abs=abs_diff, rel=rel_diff, sum=sum_diff, n=n_diff, dnan=nan_diff, argm=argm
         )
+        if dev is not None:
+            res["dev"] = dev
         if hist:
             if isinstance(hist, bool):
                 hist = np.array([0, 0.0001, 0.001, 0.01, 0.1, 1, 10, 100], dtype=diff.dtype)
@@ -1350,9 +1376,14 @@ def max_diff(
     if isinstance(expected, torch.Tensor) and isinstance(got, torch.Tensor):
         if verbose >= 6:
             print(f"[max_diff] tensor: {string_type(expected)} ? {string_type(got)}")
+        dev = 0 if expected.device == got.device else 1
         if _index < begin or (end != -1 and _index >= end):
             # out of boundary
-            return dict(abs=0.0, rel=0.0, sum=0.0, n=0.0, dnan=0)
+            if verbose >= 10:
+                if debug_info:
+                    print("\n".join(debug_info))
+                print("[max_diff] out of boundary")
+            return dict(abs=0.0, rel=0.0, sum=0.0, n=0.0, dnan=0, dev=dev)
         if expected.dtype in (torch.complex64, torch.complex128):
             if got.dtype == expected.dtype:
                 got = torch.view_as_real(got)
@@ -1446,31 +1477,63 @@ def max_diff(
                 )
 
         res: Dict[str, float] = dict(  # type: ignore
-            abs=abs_diff, rel=rel_diff, sum=sum_diff, n=n_diff, dnan=nan_diff, argm=argm
+            abs=abs_diff,
+            rel=rel_diff,
+            sum=sum_diff,
+            n=n_diff,
+            dnan=nan_diff,
+            argm=argm,
+            dev=dev,
         )
         if hist:
-            if isinstance(hist, bool):
-                hist = torch.tensor(
-                    [0, 0.0001, 0.001, 0.01, 0.1, 1, 10, 100], dtype=diff.dtype
-                )
-            hist = hist.to(diff.device)
-            ind = torch.bucketize(diff.reshape((-1,)), hist, right=False)
-            cou = torch.bincount(ind, minlength=ind.shape[0] + 1)
-            res["rep"] = dict(
-                zip(
-                    [f">{x}" for x in hist],
-                    [int(i) for i in (cou.sum() - torch.cumsum(cou, 0))],
+            if isinstance(hist, list) and len(hist) == 1:
+                res["rep"] = {f">{hist[0]}": (diff > hist[0]).sum().item()}
+            elif isinstance(hist, list) and len(hist) == 2:
+                res["rep"] = {
+                    f">{hist[0]}": (diff > hist[0]).sum().item(),
+                    f">{hist[1]}": (diff > hist[1]).sum().item(),
+                }
+            else:
+                if isinstance(hist, bool):
+                    hist = torch.tensor(
+                        [0, 0.0001, 0.001, 0.01, 0.1, 1, 10, 100], dtype=diff.dtype
+                    )
+                hist = torch.tensor(hist).to(diff.device)
+                ind = torch.bucketize(diff.reshape((-1,)), hist, right=False)
+                cou = torch.bincount(ind, minlength=ind.shape[0] + 1)
+                res["rep"] = dict(
+                    zip(
+                        [f">{x}" for x in hist],
+                        [int(i) for i in (cou.sum() - torch.cumsum(cou, 0))],
+                    )
                 )
-            )
         return res  # type: ignore
 
+    if isinstance(expected, int) and isinstance(got, torch.Tensor):
+        # a size
+        if verbose >= 6:
+            print(f"[max_diff] int: {string_type(expected)} ? {string_type(got)}")
+        if got.shape != tuple():
+            return dict(  # type: ignore
+                abs=np.inf,
+                rel=np.inf,
+                sum=np.inf,
+                n=np.inf,
+                dnan=np.inf,
+                argm=np.inf,
+            )
+        return dict(  # type: ignore
+            abs=abs(expected - got.item()),
+            rel=abs((expected - got.item()) / max(1, expected)),
+            sum=abs(expected - got.item()),
+            n=1,
+            dnan=0,
+        )
+
     if "SquashedNormal" in expected.__class__.__name__:
         if verbose >= 6:
             print(f"[max_diff] SquashedNormal: {string_type(expected)} ? {string_type(got)}")
-        values = (
-            expected.mean.detach().to("cpu"),
-            expected.scale.detach().to("cpu"),
-        )
+        values = (expected.mean, expected.scale)
         return max_diff(values, got, debug_info=_debug("SquashedNormal"), **_dkws)
 
     if expected.__class__ in torch.utils._pytree.SUPPORTED_NODES:
@@ -1675,7 +1738,7 @@ def max_diff(
 
     raise AssertionError(
         f"Not implemented with implemented with expected="
-        f"{string_type(expected)}, got={string_type(got)},\n"
+        f"{string_type(expected)} ({type(expected)}), got={string_type(got)},\n"
         f"level={level}"
     )
 
@@ -1683,6 +1746,9 @@ def max_diff(
 def string_diff(diff: Dict[str, Any]) -> str:
     """Renders discrepancies return by :func:`max_diff` into one string."""
     # dict(abs=, rel=, sum=, n=n_diff, dnan=)
+    if "dev" in diff:
+        ddiff = {k: v for k, v in diff.items() if k != "dev"}
+        return f"{string_diff(ddiff)}, dev={diff['dev']}"
     suffix = ""
     if "rep" in diff:
         rows = []

onnx_diagnostic/helpers/log_helper.py

@@ -901,13 +901,19 @@ class CubeLogs:
             else g.groupby([*key_index, *key_columns], dropna=False).sum()
         )
         not_unique = r[r["count"] > 1]
+        if not_unique.shape[0] > 0 and os.environ.get("DUPLICATE", ""):
+            filename = os.environ.get("DUPLICATE")
+            subset = data.set_index([*key_index, *key_columns]).merge(
+                not_unique.head(), left_index=True, right_index=True
+            )
+            subset.to_excel(filename)
         assert not_unique.shape[0] == 0, (
             f"view_def.name={view_def.name!r}, "
             f"unable to run the pivot with index={sorted(key_index)}, "
             f"key={sorted(key_columns)}, key_agg={key_agg}, values={sorted(values)}, "
             f"columns={sorted(data.columns)}, ignored={view_def.ignore_columns}, "
-            f"not unique={set(data.columns) - unique}"
-            f"\n--\n{not_unique.head(10)}"
+            f"not unique={set(data.columns) - unique}, set DUPLICATE=<filename> "
+            f"to store the duplicates in a excel file\n--\n{not_unique.head(10)}"
         )
 
         # pivot
@@ -1000,8 +1006,12 @@ class CubeLogs:
         keys = set(self.keys_time) - {columns_to_fix}
         select = data[self.keys_time]
         select_agg = select.groupby(list(keys)).count()
+        if select_agg.shape[0] == 0:
+            # nothing to fix
+            return data
         assert select_agg[columns_to_fix].max() <= 1, (
-            f"Column {columns_to_fix!r} has two distinct values at least for one date\n"
+            f"Column {columns_to_fix!r} has two distinct values at least for one date, "
+            f"max={select_agg[columns_to_fix].max()}\n"
             f"{select_agg[select_agg[columns_to_fix] > 1]}"
         )
 
@@ -1038,6 +1048,16 @@ class CubeLogs:
             f"data.columns.equals(res.columns)={data.columns.equals(res.columns)}, "
             f"data.index.equals(res.columns)={data.index.equals(res.columns)}, "
         )
+        select = res[self.keys_time]
+        select_agg = select.groupby(list(keys)).count()
+        if select_agg.shape[0] == 0:
+            # nothing to fix
+            return data
+        # assert select_agg[columns_to_fix].max() <= 1, (
+        #    f"Column {columns_to_fix!r} has two distinct values at least for one date, "
+        #    f"max={select_agg[columns_to_fix].max()}\n"
+        #    f"{select_agg[select_agg[columns_to_fix] > 1]}"
+        # )
         return res
 
     def _dropna(
@@ -1977,7 +1997,9 @@ class CubeLogsPerformance(CubeLogs):
         * **cmd:** command lines
         * **raw-short:** raw data without all the unused columns
         """
-        fix_aggregation_change = ["model_speedup_input_set", "model_test_with"]
+        # This does not work.
+        # used to be ["model_speedup_input_set", "model_test_with"]
+        fix_aggregation_change = []  # type: ignore[var-annotated]
         fs = ["suite", "model_suite", "task", "model_name", "model_task"]
         index_cols = self._filter_column(fs, self.keys_time)
         assert index_cols, (

onnx_diagnostic/helpers/mini_onnx_builder.py

@@ -159,7 +159,9 @@ class MiniOnnxBuilder:
         """
         if not tensors:
             # empty list
-            self.nodes.append(oh.make_node("SequenceEmpty", [], [name]))
+            self.nodes.append(
+                oh.make_node("SequenceEmpty", [], [name], dtype=TensorProto.FLOAT)
+            )
             tensor_type_proto = oh.make_tensor_type_proto(
                 elem_type=TensorProto.FLOAT, shape=None
             )
@@ -422,6 +424,27 @@ def create_onnx_model_from_input_tensors(
     :return: ModelProto
 
     The function raises an error if not supported.
+    An example:
+
+    .. code-block:: python
+
+        from onnx_diagnostic.helpers.mini_onnx_builder import (
+            create_onnx_model_from_input_tensors,
+        )
+        import onnx
+
+        proto = create_onnx_model_from_input_tensors(
+            dict(
+                query_states=query_states,
+                key_states=key_states,
+                value_states=value_states,
+                cu_seqlens=cu_seqlens,
+                max_seqlen=(cu_seqlens[1:] - cu_seqlens[:-1]).max(),
+                scaling=self.scaling,
+                attn_output=attn_output,
+            )
+        )
+        onnx.save(proto, "attention_inputs.onnx")
     """
     if switch_low_high is None:
         switch_low_high = sys.byteorder != "big"
@@ -461,7 +484,17 @@ def _unflatten(
         if spl[-1] == "array":
             return pos + 1, outputs[pos]
         if spl[-1] == "tensor":
-            return pos + 1, torch.from_numpy(outputs[pos]).to(device)
+            try:
+                return pos + 1, torch.from_numpy(outputs[pos]).to(device)
+            except TypeError:
+                # it should be more robust
+                import ml_dtypes
+
+                if outputs[pos].dtype == ml_dtypes.bfloat16:
+                    return pos + 1, torch.from_numpy(outputs[pos].astype(float)).to(device).to(
+                        torch.bfloat16
+                    )
+                raise
         raise AssertionError(f"Unexpected name {name!r} in {names}")
 
     res: List[Any] = []
@@ -532,6 +565,12 @@ def _unflatten(
             return d
         return ty(res)
 
+    if end and len(res) == 1:
+        if res[0] is None:
+            return next_pos, ty()
+        if isinstance(res[0], tuple) and len(res[0]) == 2 and res[0] == ("dict.", None):
+            return next_pos, ty()
+        return next_pos, _make(ty, res)
     return next_pos, (
         ty() if len(res) == 1 and res[0] in (("dict.", None), None) else _make(ty, res)
     )
@@ -557,6 +596,19 @@ def create_input_tensors_from_onnx_model(
     :return: restored data
 
     See example :ref:`l-plot-intermediate-results` for an example.
+
+    .. code-block:: python
+
+        import os
+        from onnx_diagnostic.helpers.mini_onnx_builder import (
+            create_input_tensors_from_onnx_model,
+        )
+        from onnx_diagnostic.helpers import string_type
+
+        restored = create_input_tensors_from_onnx_model("attention_inputs.onnx")
+        for k, v in restored.items():
+            print(f"{k}: {string_type(v, with_shape=True, with_min_max=True)}")
+
     """
     if engine == "ExtendedReferenceEvaluator":
         from ..reference import ExtendedReferenceEvaluator
@@ -595,6 +647,8 @@ def create_input_tensors_from_onnx_model(
             return float(output[0])
         if name == "tensor":
             return torch.from_numpy(output).to(device)
-        raise AssertionError(f"Unexpected name {name!r} in {names}")
+        assert name.startswith(
+            ("list_", "list.", "dict.", "tuple_", "tuple.")
+        ), f"Unexpected name {name!r} in {names}"
 
     return _unflatten(sep, names, got, device=device)[1]

onnx_diagnostic/helpers/model_builder_helper.py

@@ -28,10 +28,37 @@ def download_model_builder_to_cache(
     if file_path.exists():
         return file_path
 
+    builders = cache_dir / "builders"
+    if not builders.exists():
+        builders.mkdir(parents=True, exist_ok=True)
+
+    for subfile in [
+        "__init__.py",
+        "base.py",
+        "chatglm.py",
+        "ernie.py",
+        "gemma.py",
+        "gptoss.py",
+        "granite.py",
+        "llama.py",
+        "mistral.py",
+        "nemotron.py",
+        "olmo.py",
+        "phi.py",
+        "qwen.py",
+        "smollm.py",
+    ]:
+        u = f"{'/'.join(url.split('/')[:-1])}/builders/{subfile}"
+        response = requests.get(u)
+        response.raise_for_status()
+        with open(builders / subfile, "wb") as f:
+            f.write(response.content)
+
     response = requests.get(url)
     response.raise_for_status()
     with open(file_path, "wb") as f:
         f.write(response.content)
+
     return file_path