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

onnx_diagnostic/export/onnx_plug.py

@@ -0,0 +1,396 @@
+import inspect
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+import onnx
+import torch
+from ..helpers import max_diff
+from ..helpers.torch_helper import torch_dtype_to_onnx_dtype
+from ..reference import OnnxruntimeEvaluator
+
+TUPLE_TENSORS = Tuple[torch.Tensor, ...]
+
+
+def is_exporting() -> bool:
+    """
+    Returns :func:`torch.compiler.is_exporting` or
+    :func:`torch.compiler.is_compiling`.
+    Changes ``_TEST_EXPORT`` to make it trigger.
+    """
+    return torch.compiler.is_exporting() or torch.compiler.is_compiling()
+
+
+@dataclass
+class VerifyResult:
+    """
+    Outputs of method :meth:`verify
+    <onnx_diagnostic.export.onnx_plug.EagerDirectReplacementWithOnnx.verify>`.
+    """
+
+    eager_outputs: TUPLE_TENSORS
+    onnx_outputs: TUPLE_TENSORS
+    diffs: Tuple[Dict[str, float], ...]
+
+
+class EagerDirectReplacementWithOnnx:
+    """
+    Replaces a piece of code by another one written in ONNX
+    at export time. The function inserts a custom operator
+    and links it to the eager_fn
+
+    :param eager_fn: the code it replaces, it must be given in order to be able
+        to execute the torch.fx.Graph the exporter produces
+    :param shape_fn: the function produces dummy outputs with the shapes
+        the exporter can use for the next operators in the graph
+    :param function_proto: instances of ``onnx.FunctionProto``,
+        its domain must be ``onnx_plug``
+    :param n_inputs: number of inputs of the function, if not given,
+        the class will infer it from eager_fn signature,
+        only tensors must be counted
+    :param n_outputs: same for the number of outputs,
+        only tensors must be counted
+    :param name: the name of the custom op, the function name if not specified
+    :param kwargs: constants parameters with their default values
+    :param verbose: verbose level
+
+    Here is an example:
+
+    .. runpython::
+        :showcode:
+
+        import onnx.helper as oh
+        import torch
+        from onnx_diagnostic.helpers.onnx_helper import pretty_onnx
+        from onnx_diagnostic.export.onnx_plug import EagerDirectReplacementWithOnnx
+        from onnx_diagnostic.export.api import to_onnx
+        from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
+
+
+        def demo_customsub(x, y):
+            return x - y
+
+
+        def demo_customsub_shape(x, y):
+            return torch.empty(torch.broadcast_shapes(x.shape, y.shape), dtype=x.dtype)
+
+
+        def make_function_proto():
+            return oh.make_function(
+                "onnx_plug",
+                "demo_customsub",
+                ["x", "y"],
+                ["z"],
+                [oh.make_node("Sub", ["x", "y"], ["z"])],
+                opset_imports=[oh.make_opsetid("", 22)],
+            )
+
+
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                y = x.sum(axis=1, keepdim=True)
+                d = torch.ops.onnx_plug.demo_customsub(x, y)
+                return torch.abs(d)
+
+
+        replacements = [
+            EagerDirectReplacementWithOnnx(
+                demo_customsub, demo_customsub_shape, make_function_proto(), 2, 1
+            )
+        ]
+
+        x = torch.randn((3, 4), dtype=torch.float32)
+        model = Model()
+        ds = ({0: "d1", 1: "d2"},)
+
+        # The exported program shows a custom op.
+        ep = torch.export.export(model, (x,), dynamic_shapes=use_dyn_not_str(ds))
+        print("ep")
+
+        # As the exporter knows how the replace this custom op.
+        # Let's export.
+
+        onx = to_onnx(
+            model,
+            (x,),
+            dynamic_shapes=ds,
+            exporter="custom",
+            onnx_plugs=replacements,
+            target_opset=22,
+            inline=False,
+        ).model_proto
+
+        print(pretty_onnx(onx))
+
+        # And with :func:`torch.onnx.export`:
+
+        onx = to_onnx(
+            model,
+            (x,),
+            dynamic_shapes=ds,
+            exporter="onnx-dynamo",
+            onnx_plugs=replacements,
+            target_opset=22,
+            inline=False,
+        ).model_proto
+
+        print(pretty_onnx(onx))
+    """
+
+    def __init__(
+        self,
+        eager_fn: Callable[[TUPLE_TENSORS], TUPLE_TENSORS],
+        shape_fn: Callable[[TUPLE_TENSORS], TUPLE_TENSORS],
+        function_proto: onnx.FunctionProto,
+        n_inputs: Optional[int] = None,
+        n_outputs: Optional[int] = None,
+        name: Optional[str] = None,
+        kwargs: Optional[Dict[str, Union[int, float]]] = None,
+        verbose: int = 0,
+    ):
+        assert isinstance(
+            function_proto, onnx.FunctionProto
+        ), f"Unexpected type {type(function_proto)} for function_proto"
+        assert isinstance(n_inputs, int), f"not implemented yet when n_inputs={n_inputs}"
+        assert isinstance(n_outputs, int), f"not implemented yet when n_inputs={n_outputs}"
+        self.eager_fn = eager_fn
+        self.shape_fn = shape_fn
+        self.function_proto = function_proto
+        self.n_inputs = n_inputs
+        self.n_outputs = n_outputs
+        self.name = name or (
+            eager_fn.__name__
+            if "<" not in eager_fn.__name__
+            else eager_fn.__qualname__.replace("<locals>", "L")
+            .replace("<lambda>", "l")
+            .replace(".", "_")
+        )
+        self.kwargs = kwargs or {}
+        assert all(isinstance(v, (int, float)) for v in self.kwargs.values()), (
+            f"Only int or floats are allowed for kwargs={kwargs}, one of them "
+            f"does not respect that constraint."
+        )
+        sig = inspect.signature(self.eager_fn)
+        params = list(sig.parameters)
+        assert (
+            len(params) >= n_inputs
+        ), f"{self.eager_fn} accepts {params} as parameters < n_inputs={n_inputs}"
+        assert n_inputs == len(function_proto.input), (
+            f"Input mismatch n_inputs={n_inputs} but "
+            f"function_proto.input={function_proto.input}"
+        )
+        assert n_outputs == len(function_proto.output), (
+            f"Output mismatch n_outputs={n_outputs} but "
+            f"function_proto.output={function_proto.output}"
+        )
+        assert (
+            function_proto.domain == self.domain
+        ), f"Function domain must be {self.domain!r} but it is {function_proto.domain!r}"
+        self.args_name = [p for p in params if p not in self.kwargs]
+        self.kwargs_name = [p for p in params if p in self.kwargs]
+        self.verbose = verbose
+        self.custom_op = self._register()
+
+    @property
+    def domain(self) -> str:
+        "Returns the onnx domain."
+        return "onnx_plug"
+
+    @property
+    def target_name(self) -> str:
+        "Returns the target name (see in the exported program)."
+        return f"{self.domain}::{self.name}"
+
+    @property
+    def torch_op(self) -> Callable:
+        "Returns ``torch.ops.onny_plug.<name>``."
+        return getattr(getattr(torch.ops, self.domain), self.name).default
+
+    def __call__(self, *args, **kwargs):
+        """Calls eager_fn or shape_fn if the model is being exported."""
+        if is_exporting():
+            return self.torch_op(*args)
+        return self.eager_fn(*args, **kwargs)
+
+    def _register(self):
+        """Registers the custom op."""
+        input_args = [f"Tensor {p}" for p in self.args_name]
+        for p in self.kwargs_name:
+            val = self.kwargs[p]
+            if isinstance(val, int):
+                input_args.append(f"int {p}={val}")
+            elif isinstance(val, float):
+                input_args.append(f"float {p}={val}")
+            else:
+                raise NotImplementedError(
+                    f"kwargs {p!r} has a default value of unsupported type {type(val)}"
+                )
+
+        inputs = ", ".join(input_args)
+        schema = f"({inputs}) -> Tensor"
+        if self.n_outputs > 1:
+            schema += "[]"
+        if self.verbose:
+            print(
+                f"[EagerDirectReplacementWithOnnx._register] "
+                f"'torch.ops.{self.domain}.{self.name}"
+            )
+            print(f"[EagerDirectReplacementWithOnnx._register] schema={schema}")
+        custom_def = torch.library.CustomOpDef(self.domain, self.name, schema, self.eager_fn)
+        custom_def.register_kernel(None)(self.eager_fn)
+        custom_def._abstract_fn = self.shape_fn
+
+    def verify(
+        self,
+        *args,
+        engine: Optional[Callable] = None,
+        dump_onnx_model: Optional[str] = None,
+        **kwargs,
+    ) -> VerifyResult:
+        """
+        Verifies that the eager mode is equivalent to the onnx function given
+        as a replacements. This function evaluates `eager_fn`, checks that the shapes
+        are equivalent to the ones given by `shape_fn`, and finally evaluates the
+        onnx translation if the previous did not fail.
+
+        :param args: function inputs
+        :param kwargs: arguments for eager_fn
+        :param engine: by default an instance of
+            :class:`onnx_diagnostic.reference.OnnxruntimeEvaluator`.
+        :param dump_onnx_model: to dump the onnx model used to verify
+            eager and onnx produce the same results
+        :param kwargs: additional arguments to the function
+        :return: outputs of :func:`onnx_diagnostic.helpers.max_diff`
+        """
+        expected = self.eager_fn(*args, **kwargs)
+        shapes = self.shape_fn(*args, **kwargs)
+        if isinstance(expected, torch.Tensor):
+            expected = (expected,)
+            assert isinstance(shapes, torch.Tensor), (
+                f"eager_fn={self.eager_fn} returns a Tensor but shape_fn={self.shape_fn} "
+                f"returns a {type(shapes)}"
+            )
+            shapes = (shapes,)
+        assert isinstance(expected, tuple) and isinstance(shapes, tuple), (
+            f"eager_fn={self.eager_fn} returns a {type(expected)} "
+            f"and shape_fn={self.shape_fn} returns a {type(shapes)}"
+        )
+        assert len(expected) and len(shapes), (
+            f"eager_fn={self.eager_fn} and shape_fn={self.shape_fn} "
+            f"do not return the same number of tensors."
+        )
+        for i, (e, s) in enumerate(zip(expected, shapes)):
+            assert e.dtype == s.dtype, (
+                f"Type mismatch {e.dtype} != {s.dtype} for output {i}, "
+                f"eager_fn={self.eager_fn} and shape_fn={self.shape_fn}"
+            )
+            assert e.shape == s.shape, (
+                f"Type mismatch {e.shape} != {s.shape} for output {i}, "
+                f"eager_fn={self.eager_fn} and shape_fn={self.shape_fn}"
+            )
+
+        # Now the ONNX execution.
+        assert engine is None, f"Not implemented yet with engine={engine!r}"
+        ags, kws = self._make_args_kwargs(*args, **kwargs)
+        sess = OnnxruntimeEvaluator(
+            self.function_proto,
+            whole=True,
+            dump_onnx_model=dump_onnx_model,
+            function_kwargs=kws,
+        )
+        feeds = dict(zip(sess.input_names, ags))
+        got = sess.run(None, feeds)
+        diffs = tuple(max_diff(e, g, hist=[0.1, 0.01]) for e, g in zip(expected, got))
+        return VerifyResult(eager_outputs=expected, onnx_outputs=tuple(got), diffs=diffs)  # type: ignore[arg-type]
+
+    def _make_args_kwargs(self, *args, **kwargs):
+        ags = args[: len(self.args_name)]
+        kws = dict(zip(self.kwargs_name, args[len(self.args_name) :]))
+        kws.update(kwargs)
+        return ags, kws
+
+    def custom_converter(
+        self,
+    ) -> Callable:
+        """
+        Returns a function which
+        converts a custom ops found in the fx graph into ONNX
+        following the API of the custom exporter.
+        The converter adds a custom op and registers the local function.
+        """
+
+        def converter(
+            g: Any,  # GraphBuilder
+            sts: Optional[Dict[str, Any]],
+            outputs: List[str],
+            *args,
+            **kwargs,
+        ) -> Any:
+            if not g.has_local_function(
+                self.function_proto.name, domain=self.function_proto.domain
+            ):
+                g.add_function(self.function_proto)
+            ags, kws = self._make_args_kwargs(*args, **kwargs)
+            res = g.make_node(
+                self.function_proto.name,
+                ags,
+                outputs,
+                domain=self.function_proto.domain,
+                name=self.target_name,
+                **kws,
+            )
+            if not sts:
+                new_shapes = self.shape_fn(*args)
+                if not isinstance(new_shapes, tuple):
+                    new_shapes = (new_shapes,)
+                for sh, o in zip(new_shapes, outputs):
+                    g.set_type(o, torch_dtype_to_onnx_dtype(sh.dtype))
+                    g.set_shape(o, sh.shape)
+            return res
+
+        return converter
+
+    def onnx_dynamo_converter(self) -> Callable:
+        """
+        Returns a function which
+        which converts a custom ops found in the fx graph into ONNX
+        following the API of :func:`torch.onnx.export`.
+        """
+        import onnxscript
+
+        onnx_plug_op = onnxscript.values.Opset(domain=self.function_proto.domain, version=1)
+        schema = onnx_plug_op[self.function_proto.name]
+        if schema is None:
+            all_types = [
+                "tensor(float)",
+                "tensor(float16)",
+                "tensor(bfloat16)",
+                "tensor(double)",
+                "tensor(int64)",
+                "tensor(int32)",
+            ]
+            type_constraints = []
+            for i in range(self.n_inputs):
+                type_constraints.append((f"T{i}", all_types, ""))
+            for i in range(self.n_outputs):
+                type_constraints.append((f"U{i}", all_types, ""))
+            schema = onnx.defs.OpSchema(
+                self.function_proto.name,
+                self.function_proto.domain,
+                1,
+                inputs=[
+                    onnx.defs.OpSchema.FormalParameter(f"arg_{i}", f"T{i}")
+                    for i in range(self.n_inputs)
+                ],
+                outputs=[
+                    onnx.defs.OpSchema.FormalParameter(f"res_{i}", f"U{i}")
+                    for i in range(self.n_outputs)
+                ],
+                type_constraints=type_constraints,
+            )
+            onnx.defs.register_schema(schema)
+        op = onnxscript.values.Op(onnx_plug_op, self.function_proto.name, schema)
+
+        def converter(*cargs, **ckwargs):
+            ags, kws = self._make_args_kwargs(*cargs, **ckwargs)
+            return op(*ags, n_outputs=self.n_outputs, **kws)
+
+        return onnxscript.values.TracedOnnxFunction(onnx_plug_op, converter)

onnx_diagnostic/ext_test_case.py

@@ -9,6 +9,7 @@ import itertools
 import logging
 import os
 import re
+import shutil
 import sys
 import unittest
 import warnings
@@ -63,7 +64,7 @@ def skipif_ci_apple(msg) -> Callable:
     return lambda x: x
 
 
-def unit_test_going():
+def unit_test_going() -> bool:
     """
     Enables a flag telling the script is running while testing it.
     Avois unit tests to be very long.
@@ -147,7 +148,7 @@ def hide_stdout(f: Optional[Callable] = None) -> Callable:
 
     def wrapper(fct):
         def call_f(self):
-            if os.environ.get("UNHIDE", ""):
+            if os.environ.get("UNHIDE", "") in (1, "1", "True", "true"):
                 fct(self)
                 return
             st = StringIO()
@@ -742,8 +743,15 @@ class ExtTestCase(unittest.TestCase):
     _warns: List[Tuple[str, int, Warning]] = []
     _todos: List[Tuple[Callable, str]] = []
 
+    def unit_test_going(self) -> bool:
+        """
+        Enables a flag telling the script is running while testing it.
+        Avois unit tests to be very long.
+        """
+        return unit_test_going()
+
     @property
-    def verbose(self):
+    def verbose(self) -> int:
         "Returns the the value of environment variable ``VERBOSE``."
         return int(os.environ.get("VERBOSE", "0"))
 
@@ -768,13 +776,13 @@ class ExtTestCase(unittest.TestCase):
         cls._todos.append((f, msg))
 
     @classmethod
-    def ort(cls):
+    def ort(cls) -> unittest.__class__:
         import onnxruntime
 
         return onnxruntime
 
     @classmethod
-    def to_onnx(self, *args, **kwargs):
+    def to_onnx(self, *args, **kwargs) -> "ModelProto":  # noqa: F821
         from experimental_experiment.torch_interpreter import to_onnx
 
         return to_onnx(*args, **kwargs)
@@ -806,12 +814,16 @@ class ExtTestCase(unittest.TestCase):
             os.makedirs(folder)
         return folder
 
-    def dump_onnx(
-        self,
-        name: str,
-        proto: Any,
-        folder: Optional[str] = None,
-    ) -> str:
+    def clean_dump(self, folder: str = "dump_test"):
+        """Cleans this folder."""
+        for item in os.listdir(folder):
+            item_path = os.path.join(folder, item)
+            if os.path.isfile(item_path) or os.path.islink(item_path):
+                os.remove(item_path)
+            elif os.path.isdir(item_path):
+                shutil.rmtree(item_path)
+
+    def dump_onnx(self, name: str, proto: Any, folder: Optional[str] = None) -> str:
         """Dumps an onnx file."""
         fullname = self.get_dump_file(name, folder=folder)
         with open(fullname, "wb") as f:
@@ -1094,10 +1106,15 @@ class ExtTestCase(unittest.TestCase):
             value = numpy.array(value).astype(expected.dtype)
         self.assertEqualArray(expected, value, atol=atol, rtol=rtol)
 
-    def check_ort(self, onx: "onnx.ModelProto") -> bool:  # noqa: F821
+    def check_ort(
+        self, onx: "onnx.ModelProto"  # noqa: F821
+    ) -> "onnxruntime.InferenceSession":  # noqa: F821
         from onnxruntime import InferenceSession
 
-        return InferenceSession(onx.SerializeToString(), providers=["CPUExecutionProvider"])
+        return InferenceSession(
+            onx if isinstance(onx, str) else onx.SerializeToString(),
+            providers=["CPUExecutionProvider"],
+        )
 
     def assertRaise(self, fct: Callable, exc_type: type[Exception], msg: Optional[str] = None):
         """In the name"""
@@ -1137,7 +1154,7 @@ class ExtTestCase(unittest.TestCase):
         if not full.endswith(suffix):
             raise AssertionError(f"suffix={suffix!r} does not end string  {full!r}.")
 
-    def capture(self, fct: Callable):
+    def capture(self, fct: Callable) -> Tuple[Any, str, str]:
         """
         Runs a function and capture standard output and error.
 
@@ -1189,6 +1206,7 @@ class ExtTestCase(unittest.TestCase):
         expected: Optional[Any] = None,
         use_ort: bool = False,
         ort_optimized_graph: bool = False,
+        ep: Optional[Union["torch.export.ExportedProgram", str]] = None,  # noqa: F821
         **kwargs,
     ):
         """
@@ -1208,6 +1226,7 @@ class ExtTestCase(unittest.TestCase):
         :param copy_inputs: to copy the inputs
         :param use_ort: use :class:`onnxruntime.InferenceSession`
         :param ort_optimized_graph: dumps the optimized onnxruntime graph
+        :param ep: exported program (or saved exported program)
         :param kwargs: arguments sent to
             :class:`onnx_diagnostic.helpers.ort_session.InferenceSessionForTorch`
         """
@@ -1226,15 +1245,16 @@ class ExtTestCase(unittest.TestCase):
             if isinstance(proto, str):
                 name = proto
                 proto = onnx.load(name)
-            else:
+            elif not self.unit_test_going():
                 assert isinstance(
                     proto, onnx.ModelProto
                 ), f"Unexpected type {type(proto)} for proto"
                 name = self.dump_onnx(name, proto)
-            if verbose:
+            if verbose and not self.unit_test_going():
                 print(f"[{vname}] file size {os.stat(name).st_size // 2**10:1.3f} kb")
         if verbose:
             print(f"[{vname}] make feeds {string_type(inputs, **kws)}")
+
         if use_ort:
             assert isinstance(
                 proto, onnx.ModelProto
@@ -1242,15 +1262,14 @@ class ExtTestCase(unittest.TestCase):
             feeds = make_feeds(proto, inputs, use_numpy=True, copy=True)
             import onnxruntime
 
-            if verbose:
-                print(f"[{vname}] create onnxruntime.InferenceSession")
             options = onnxruntime.SessionOptions()
             if ort_optimized_graph:
                 options.optimized_model_filepath = f"{name}.optort.onnx"
+            providers = kwargs.get("providers", ["CPUExecutionProvider"])
+            if verbose:
+                print(f"[{vname}] create onnxruntime.InferenceSession with {providers}")
             sess = onnxruntime.InferenceSession(
-                proto.SerializeToString(),
-                options,
-                providers=kwargs.get("providers", ["CPUExecutionProvider"]),
+                proto.SerializeToString(), options, providers=providers
             )
             if verbose:
                 print(f"[{vname}] run ort feeds {string_type(feeds, **kws)}")
@@ -1265,6 +1284,7 @@ class ExtTestCase(unittest.TestCase):
             got = sess.run(None, feeds)
         if verbose:
             print(f"[{vname}] compute expected values")
+
         if expected is None:
             if copy_inputs:
                 expected = (
@@ -1274,10 +1294,46 @@ class ExtTestCase(unittest.TestCase):
                 )
             else:
                 expected = model(*inputs) if isinstance(inputs, tuple) else model(**inputs)
+
         if verbose:
             print(f"[{vname}] expected {string_type(expected, **kws)}")
             print(f"[{vname}] obtained {string_type(got, **kws)}")
-        diff = max_diff(expected, got, flatten=True)
+
+        if ep:
+            if isinstance(ep, str):
+                if verbose:
+                    print(f"[{vname}] load exported program {ep!r}")
+                import torch
+
+                ep = torch.export.load(ep)
+            ep_inputs = copy.deepcopy(inputs) if copy_inputs else inputs
+            ep_model = ep.module()  # type: ignore[union-attr]
+            ep_expected = (
+                ep_model(*copy.deepcopy(ep_inputs))
+                if isinstance(ep_inputs, tuple)
+                else ep_model(**copy.deepcopy(ep_inputs))
+            )
+            if verbose:
+                print(f"[{vname}] ep_expected {string_type(ep_expected, **kws)}")
+            ep_diff = max_diff(expected, ep_expected, hist=[0.1, 0.01])
+            if verbose:
+                print(f"[{vname}] ep_diff {string_diff(ep_diff)}")
+            assert (
+                isinstance(ep_diff["abs"], float)
+                and isinstance(ep_diff["rel"], float)
+                and not numpy.isnan(ep_diff["abs"])
+                and ep_diff["abs"] <= atol
+                and not numpy.isnan(ep_diff["rel"])
+                and ep_diff["rel"] <= rtol
+            ), (
+                f"discrepancies in {test_name!r} between the exported program "
+                f"and the exported model diff={string_diff(ep_diff)}"
+            )
+            ep_nx_diff = max_diff(ep_expected, got, flatten=True, hist=[0.1, 0.01])
+            if verbose:
+                print(f"[{vname}] ep_nx_diff {string_diff(ep_nx_diff)}")
+
+        diff = max_diff(expected, got, flatten=True, hist=[0.1, 0.01])
         if verbose:
             print(f"[{vname}] diff {string_diff(diff)}")
         assert (
@@ -1287,7 +1343,10 @@ class ExtTestCase(unittest.TestCase):
             and diff["abs"] <= atol
             and not numpy.isnan(diff["rel"])
             and diff["rel"] <= rtol
-        ), f"discrepancies in {test_name!r}, diff={string_diff(diff)}"
+        ), (
+            f"discrepancies in {test_name!r} between the model and "
+            f"the onnx model diff={string_diff(diff)}"
+        )
 
     def _debug(self):
         "Tells if DEBUG=1 is set up."
@@ -1298,6 +1357,16 @@ class ExtTestCase(unittest.TestCase):
 
         return string_type(*args, **kwargs)
 
+    def max_diff(self, *args, **kwargs):
+        from .helpers import max_diff
+
+        return max_diff(*args, **kwargs)
+
+    def use_dyn_not_str(self, *args, **kwargs):
+        from onnx_diagnostic.torch_export_patches.patch_inputs import use_dyn_not_str
+
+        return use_dyn_not_str(*args, *kwargs)
+
     def subloop(self, *args, verbose: int = 0):
         "Loops over elements and calls :meth:`unittests.TestCase.subTest`."
         if len(args) == 1:

onnx_diagnostic/helpers/cache_helper.py

@@ -80,7 +80,7 @@ def flatten_unflatten_for_dynamic_shapes(
     start = 0
     end = 0
     subtrees = []
-    for subspec in spec.children_specs:
+    for subspec in (spec.children() if hasattr(spec, "children") else spec.children_specs):
         end += subspec.num_leaves
         value = subspec.unflatten(flat[start:end])
         value = flatten_unflatten_for_dynamic_shapes(