onnx-diagnostic 0.6.0__py3-none-any.whl → 0.6.2__py3-none-any.whl

onnx_diagnostic/reference/torch_ops/unary_ops.py

@@ -0,0 +1,86 @@
+import torch
+from . import OpRunKernel, OpRunTensor
+
+
+class Abs_1(OpRunKernel):
+    """Abs"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(torch.abs(x.tensor))
+
+
+class Cos_1(OpRunKernel):
+    """Cos"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor.cos())
+
+
+class Erf_9(OpRunKernel):
+    """Erf"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor.erf())
+
+
+class Exp_1(OpRunKernel):
+    """Exp"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor.exp())
+
+
+class Identity_1(OpRunKernel):
+    "Identity"
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor)
+
+
+class Log_1(OpRunKernel):
+    """Log"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor.log())
+
+
+class Neg_1(OpRunKernel):
+    """Neg"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(-x.tensor)
+
+
+class Not_1(OpRunKernel):
+    """Not"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(~x.tensor)
+
+
+class Reciprocal_1(OpRunKernel):
+    """REciprocal"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(1 / x.tensor)
+
+
+class Sigmoid_6(OpRunKernel):
+    """Sqrt"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(torch.sigmoid(x.tensor))
+
+
+class Sin_1(OpRunKernel):
+    """Sin"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor.sin())
+
+
+class Sqrt_1(OpRunKernel):
+    """Sqrt"""
+
+    def run(self, x: OpRunTensor) -> OpRunTensor:
+        return OpRunTensor(x.tensor.sqrt())

onnx_diagnostic/tasks/__init__.py

@@ -39,9 +39,30 @@ def supported_tasks() -> List[str]:
 
 def reduce_model_config(config: Any, task: str) -> Dict[str, Any]:
     """Reduces a model size."""
+    head_size0 = (
+        config.head_dim
+        if hasattr(config, "head_dim") and config.head_dim
+        else (
+            config.hidden_size // config.num_attention_heads
+            if hasattr(config, "hidden_size") and hasattr(config, "num_attention_heads")
+            else None
+        )
+    )
     tasks = {mod.__TASK__: mod.reduce_model_config for mod in __TASKS__}
     assert task in tasks, f"Task {task!r} not found in {sorted(tasks)}"
-    return tasks[task](config)
+    res = tasks[task](config)
+    if head_size0 and "head_dim" in res:
+        head_size = (
+            config.head_dim
+            if hasattr(config, "head_dim") and config.head_dim
+            else config.hidden_size // config.num_attention_heads
+        )
+        assert head_size0 == head_size or head_size % 16 == 0, (
+            f"head_size should be a multiple of 16 "
+            f"(head_size0={head_size0}), res={res}, "
+            f"config=\n{config}"
+        )
+    return res
 
 
 def random_input_kwargs(config: Any, task: str) -> Tuple[Dict[str, Any], Callable]:

onnx_diagnostic/tasks/image_classification.py

@@ -58,8 +58,8 @@ def get_inputs(
     shapes = {
         "pixel_values": {
             0: torch.export.Dim("batch", min=1, max=1024),
-            2: torch.export.Dim("width", min=1, max=4096),
-            3: torch.export.Dim("height", min=1, max=4096),
+            2: "width",
+            3: "height",
         },
     }
     inputs = dict(

onnx_diagnostic/tasks/text_generation.py

@@ -27,7 +27,7 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
         kwargs = dict(
             num_hidden_layers=min(config.num_hidden_layers, 2),
             intermediate_size=256 if config is None else min(512, config.intermediate_size),
-            hidden_size=256 if config is None else min(256, config.hidden_size),
+            hidden_size=512 if config is None else min(512, config.hidden_size),
             cls_cache="MambaCache",
             state_size=8 if config is None else getattr(config, "state_size", None),
             conv_kernel=4 if config is None else getattr(config, "conv_kernel", None),
@@ -44,8 +44,8 @@ def reduce_model_config(config: Any) -> Dict[str, Any]:
                 else config.num_attention_heads
             ),
             hidden_size=(
-                min(config.hidden_size, 3072 // 4)
-                if config.hidden_size % 4 == 0
+                min(config.hidden_size, 4096 // 4)
+                if config.hidden_size % 64 == 0
                 else config.hidden_size
             ),
         )

onnx_diagnostic/torch_export_patches/eval/__init__.py

@@ -185,9 +185,17 @@ def _make_exporter_export(
 
     if exporter == "export-strict":
         try:
-            exported = torch.export.export(
-                model, inputs, dynamic_shapes=dynamic_shapes, strict=True
-            )
+            if verbose >= 2:
+                exported = torch.export.export(
+                    model, inputs, dynamic_shapes=dynamic_shapes, strict=True
+                )
+            else:
+                with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(
+                    io.StringIO()
+                ):
+                    exported = torch.export.export(
+                        model, inputs, dynamic_shapes=dynamic_shapes, strict=True
+                    )
         except Exception as e:
             if not quiet:
                 raise
@@ -198,17 +206,33 @@ def _make_exporter_export(
         return exported.module()
     if exporter in ("export-strict-dec", "export-strict-decall"):
         try:
-            exported = torch.export.export(
-                model, inputs, dynamic_shapes=dynamic_shapes, strict=True
-            )
-            if verbose >= 9:
-                print("-- graph before decomposition")
-                print(exported.graph)
-            exported = (
-                exported.run_decompositions()
-                if "decall" in exporter
-                else exported.run_decompositions({})
-            )
+            if verbose >= 2:
+                exported = torch.export.export(
+                    model, inputs, dynamic_shapes=dynamic_shapes, strict=True
+                )
+                if verbose >= 9:
+                    print("-- graph before decomposition")
+                    print(exported.graph)
+                exported = (
+                    exported.run_decompositions()
+                    if "decall" in exporter
+                    else exported.run_decompositions({})
+                )
+            else:
+                with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(
+                    io.StringIO()
+                ):
+                    exported = torch.export.export(
+                        model, inputs, dynamic_shapes=dynamic_shapes, strict=True
+                    )
+                    if verbose >= 9:
+                        print("-- graph before decomposition")
+                        print(exported.graph)
+                    exported = (
+                        exported.run_decompositions()
+                        if "decall" in exporter
+                        else exported.run_decompositions({})
+                    )
         except Exception as e:
             if not quiet:
                 raise
@@ -219,9 +243,17 @@ def _make_exporter_export(
         return exported.module()
     if exporter == "export-nostrict":
         try:
-            exported = torch.export.export(
-                model, inputs, dynamic_shapes=dynamic_shapes, strict=False
-            )
+            if verbose >= 2:
+                exported = torch.export.export(
+                    model, inputs, dynamic_shapes=dynamic_shapes, strict=False
+                )
+            else:
+                with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(
+                    io.StringIO()
+                ):
+                    exported = torch.export.export(
+                        model, inputs, dynamic_shapes=dynamic_shapes, strict=False
+                    )
         except Exception as e:
             if not quiet:
                 raise
@@ -232,17 +264,33 @@ def _make_exporter_export(
         return exported.module()
     if exporter in ("export-nostrict-dec", "export-nostrict-decall"):
         try:
-            exported = torch.export.export(
-                model, inputs, dynamic_shapes=dynamic_shapes, strict=False
-            )
-            if verbose >= 9:
-                print("-- graph before decomposition")
-                print(exported.graph)
-            exported = (
-                exported.run_decompositions()
-                if "decall" in exporter
-                else exported.run_decompositions({})
-            )
+            if verbose >= 2:
+                exported = torch.export.export(
+                    model, inputs, dynamic_shapes=dynamic_shapes, strict=False
+                )
+                if verbose >= 9:
+                    print("-- graph before decomposition")
+                    print(exported.graph)
+                exported = (
+                    exported.run_decompositions()
+                    if "decall" in exporter
+                    else exported.run_decompositions({})
+                )
+            else:
+                with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(
+                    io.StringIO()
+                ):
+                    exported = torch.export.export(
+                        model, inputs, dynamic_shapes=dynamic_shapes, strict=False
+                    )
+                    if verbose >= 9:
+                        print("-- graph before decomposition")
+                        print(exported.graph)
+                    exported = (
+                        exported.run_decompositions()
+                        if "decall" in exporter
+                        else exported.run_decompositions({})
+                    )
         except Exception as e:
             if not quiet:
                 raise
@@ -255,8 +303,15 @@ def _make_exporter_export(
         from experimental_experiment.torch_interpreter.tracing import CustomTracer
 
         try:
-            graph = CustomTracer().trace(model)
-            mod = torch.fx.GraphModule(model, graph)
+            if verbose >= 2:
+                graph = CustomTracer().trace(model)
+                mod = torch.fx.GraphModule(model, graph)
+            else:
+                with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(
+                    io.StringIO()
+                ):
+                    graph = CustomTracer().trace(model)
+                    mod = torch.fx.GraphModule(model, graph)
         except Exception as e:
             if not quiet:
                 raise
@@ -289,13 +344,25 @@ def _make_exporter_onnx(
         if "-dec" in exporter:
             opts["decomposition_table"] = "all" if "-decall" in exporter else "default"
         try:
-            onx, builder = to_onnx(
-                model,
-                inputs,
-                dynamic_shapes=dynamic_shapes,
-                export_options=ExportOptions(**opts),
-                return_builder=True,
-            )
+            if verbose >= 2:
+                onx, builder = to_onnx(
+                    model,
+                    inputs,
+                    dynamic_shapes=dynamic_shapes,
+                    export_options=ExportOptions(**opts),
+                    return_builder=True,
+                )
+            else:
+                with contextlib.redirect_stdout(io.StringIO()), contextlib.redirect_stderr(
+                    io.StringIO()
+                ):
+                    onx, builder = to_onnx(
+                        model,
+                        inputs,
+                        dynamic_shapes=dynamic_shapes,
+                        export_options=ExportOptions(**opts),
+                        return_builder=True,
+                    )
         except Exception as e:
             if not quiet:
                 raise RuntimeError(
@@ -306,6 +373,7 @@ def _make_exporter_onnx(
                 ) from e
             return dict(error=str(e), success=0, error_step="export")
         return onx, builder
+
     if exporter == "dynamo":
         import torch
 
@@ -338,6 +406,7 @@ def _make_exporter_onnx(
                 ) from e
             return dict(error=str(e), success=0, error_step="export")
         return onx, None
+
     if exporter == "dynamo-ir":
         import torch
 

onnx_diagnostic/torch_export_patches/eval/model_cases.py

@@ -23,7 +23,6 @@ class AtenRollPos(torch.nn.Module):
 
 
 class InplaceAdd(torch.nn.Module):
-
     def __init__(self):
         super().__init__()
         self.bias = torch.ones((1, 4), dtype=torch.float32)
@@ -37,7 +36,6 @@ class InplaceAdd(torch.nn.Module):
 
 
 class InplaceAdd2(torch.nn.Module):
-
     def __init__(self):
         super().__init__()
         self.bias = torch.ones((1, 4), dtype=torch.float32)
@@ -51,7 +49,6 @@ class InplaceAdd2(torch.nn.Module):
 
 
 class InplaceAdd_Mul(torch.nn.Module):
-
     def __init__(self):
         super().__init__()
         self.bias = torch.ones((1, 4), dtype=torch.float32)
@@ -65,7 +62,6 @@ class InplaceAdd_Mul(torch.nn.Module):
 
 
 class InplaceCloneAdd_(torch.nn.Module):
-
     def __init__(self):
         super().__init__()
         self.bias = torch.ones((1, 4), dtype=torch.float32)
@@ -80,7 +76,6 @@ class InplaceCloneAdd_(torch.nn.Module):
 
 
 class InplaceSetItemSquare(torch.nn.Module):
-
     def forward(self, x):
         x[:2, :3] = 1
         return x
@@ -90,7 +85,6 @@ class InplaceSetItemSquare(torch.nn.Module):
 
 
 class InplaceSetItemSquareAdd(torch.nn.Module):
-
     def forward(self, x):
         x[:2, :3] = 1
         return x + 2
@@ -100,7 +94,6 @@ class InplaceSetItemSquareAdd(torch.nn.Module):
 
 
 class InplaceSetItemSquareAdd2(torch.nn.Module):
-
     def forward(self, x):
         x[:2, :3] = 1
         return x + 2, x + 3
@@ -110,7 +103,6 @@ class InplaceSetItemSquareAdd2(torch.nn.Module):
 
 
 class InplaceSetItemEllipsis_1(torch.nn.Module):
-
     def __init__(self):
         super().__init__()
         self.params = torch.zeros((1, 8192, 4), dtype=torch.float32)
@@ -124,11 +116,10 @@ class InplaceSetItemEllipsis_1(torch.nn.Module):
         (torch.from_numpy(np.array([0, 3, 2, 1])).to(torch.int64)),
         (torch.arange(4 * 8192) + 10).reshape((-1, 4)).to(torch.float32),
     )
-    _dynamic = {"x": {0: DIM("batch")}}
+    _dynamic = {"index": {0: DIM("batch")}, "update": {0: DIM("batch"), 1: DYN}}
 
 
 class InplaceSetItemEllipsis_2(torch.nn.Module):
-
     def __init__(self):
         super().__init__()
         self.params = torch.zeros((1, 8192, 6), dtype=torch.float32)
@@ -142,7 +133,7 @@ class InplaceSetItemEllipsis_2(torch.nn.Module):
         torch.from_numpy(np.array([0, 3, 2, 5])).to(torch.int64),
         (torch.arange(4 * 8192) + 10).reshape((-1, 4)).to(torch.float32),
     )
-    _dynamic = {"x": {0: DIM("batch")}}
+    _dynamic = {"index": {0: DIM("batch")}, "update": {0: DIM("batch"), 1: DYN}}
 
 
 class InplaceSetItemMask(torch.nn.Module):
@@ -156,7 +147,6 @@ class InplaceSetItemMask(torch.nn.Module):
 
 
 class AtenInterpolate(torch.nn.Module):
-
     def forward(self, x):
         y = torch.nn.functional.interpolate(
             x,
@@ -171,7 +161,6 @@ class AtenInterpolate(torch.nn.Module):
 
 
 class AtenNonZero(torch.nn.Module):
-
     def forward(self, x):
         y = torch.nonzero(x)
         return y
@@ -181,7 +170,6 @@ class AtenNonZero(torch.nn.Module):
 
 
 class AtenNonZeroTuple(torch.nn.Module):
-
     def forward(self, x):
         y = torch.nonzero(x, as_tuple=True)
         return y[0], y[1]
@@ -191,7 +179,6 @@ class AtenNonZeroTuple(torch.nn.Module):
 
 
 class AtenAsStrided(torch.nn.Module):
-
     def __init__(self):
         super().__init__()
 
@@ -288,7 +275,6 @@ class ControlFlowCondConstant(torch.nn.Module):
 
 
 class ControlFlowCondNestedModule(torch.nn.Module):
-
     class Submodule(torch.nn.Module):
         def __init__(self):
             super().__init__()
@@ -367,7 +353,9 @@ class ControlFlowCondNonZero(torch.nn.Module):
 
 
 class ControlFlowCondIdentity_153832(torch.nn.Module):
-    """`#153832 <https://github.com/pytorch/pytorch/issues/153832>`_"""
+    """
+    `#153832 <https://github.com/pytorch/pytorch/issues/153832>`_
+    """
 
     def forward(self, x, y):
 
@@ -389,7 +377,6 @@ class ControlFlowCondIdentity_153832(torch.nn.Module):
 
 
 class ControlFlowScan(torch.nn.Module):
-
     @staticmethod
     def add(carry: torch.Tensor, y: torch.Tensor):
         next_carry = carry + y
@@ -486,7 +473,6 @@ class ControlFlowScanCDist2(torch.nn.Module):
 
 
 class ControlFlowScanCDistXY(torch.nn.Module):
-
     @staticmethod
     def dist(y: torch.Tensor, scanned_x: torch.Tensor):
         sub = y - scanned_x.reshape((1, -1))
@@ -517,7 +503,9 @@ class ControlFlowScanCDistXY(torch.nn.Module):
 
 
 class ControlFlowScanInplace_153705(torch.nn.Module):
-    """`#153705 <https://github.com/pytorch/pytorch/issues/153705>`_"""
+    """
+    `#153705 <https://github.com/pytorch/pytorch/issues/153705>`_
+    """
 
     def forward(self, x, y):
         def loop_body_1(z, iv, x, y):
@@ -540,7 +528,9 @@ class ControlFlowScanInplace_153705(torch.nn.Module):
 
 
 class ControlFlowScanDecomposition_151564(torch.nn.Module):
-    """`#151564 <https://github.com/pytorch/pytorch/issues/151564>`_"""
+    """
+    `#151564 <https://github.com/pytorch/pytorch/issues/151564>`_
+    """
 
     @classmethod
     def dummy_loop(cls, padded: torch.Tensor, pos: torch.Tensor):
@@ -790,7 +780,6 @@ class SignatureShapeAsIndex(torch.nn.Module):
 
 
 class TypeBFloat16(torch.nn.Module):
-
     def forward(self, x):
         xb = x.to(torch.bfloat16)
         return (xb + xb).to(torch.float32)
@@ -825,7 +814,6 @@ class CropLastDimensionWithTensorShape(torch.nn.Module):
 
 
 class CropLastDimensionWithTensorContent(torch.nn.Module):
-
     def forward(self, x, shape):
         return x[..., : shape[0]]
 
@@ -833,11 +821,10 @@ class CropLastDimensionWithTensorContent(torch.nn.Module):
         (torch.rand(3, 4, 4).to(torch.float32), torch.tensor([2], dtype=torch.int64)),
         (torch.rand(6, 4, 4).to(torch.float32), torch.tensor([3], dtype=torch.int64)),
     ]
-    _dynamic = {"x": {0: DIM("batch")}}
+    _dynamic = {"x": {0: DIM("batch")}, "shape": {}}
 
 
 class SignatureListFixedWithNone(torch.nn.Module):
-
     def forward(self, lx):
         x = lx[0]
         if lx[1] is not None: