tico 0.1.0.dev250520__py3-none-any.whl → 0.1.0.dev250521__py3-none-any.whl

tico/__init__.py

@@ -22,7 +22,7 @@ from tico.config import CompileConfigV1, get_default_config
 from tico.utils.convert import convert, convert_from_exported_program, convert_from_pt2
 
 # THIS LINE IS AUTOMATICALLY GENERATED BY setup.py
-__version__ = "0.1.0.dev250520"
+__version__ = "0.1.0.dev250521"
 
 
 if Version(torch.__version__) < Version("2.5.0"):

tico/passes/decompose_group_norm.py

@@ -88,6 +88,25 @@ class DecomposeGroupNorm(PassBase):
     def __init__(self):
         super().__init__()
 
+    def _insert_norm(self, graph, tensor, eps):
+        """
+        Insert (tensor - mean) / sqrt(var + eps)) into the graph
+          and return the normalized tensor node.
+        """
+        mean = graph.call_function(
+            torch.ops.aten.mean.dim, (tensor, [-1]), {"keepdim": True}
+        )
+        deviation = graph.call_function(torch.ops.aten.sub.Tensor, (tensor, mean))
+        squared = graph.call_function(torch.ops.aten.pow.Tensor_Scalar, (deviation, 2))
+        var = graph.call_function(
+            torch.ops.aten.mean.dim, (squared, [-1]), {"keepdim": True}
+        )
+        inverse_std = graph.call_function(
+            torch.ops.aten.rsqrt.default,
+            (graph.call_function(torch.ops.aten.add.Tensor, (var, eps)),),
+        )
+        return graph.call_function(torch.ops.aten.mul.Tensor, (deviation, inverse_std))
+
     def call(self, exported_program: ExportedProgram) -> PassResult:
         logger = logging.getLogger(__name__)
 
@@ -155,52 +174,20 @@ class DecomposeGroupNorm(PassBase):
             pack_shape = [layer_size, norm_size]
 
             with gm.graph.inserting_before(node):
-                layer = graph.call_function(
-                    # Sometimes, `x` has a stride for NHWC, which can't be reshaped with `aten.view.default`.
-                    # TODO Find out how to process such case properly.
-                    torch.ops.aten.reshape.default,
-                    (x, pack_shape),
-                )
-                layer_mean = graph.call_function(
-                    torch.ops.aten.mean.dim,
-                    (layer, [-1]),
-                )
-                layer_mean_reshape = graph.call_function(
-                    torch.ops.aten.view.default,
-                    (layer_mean, [layer_size, 1]),
-                )
-                layer_deviation = graph.call_function(
-                    torch.ops.aten.sub.Tensor,
-                    (layer, layer_mean_reshape),
-                )
-                layer_sqr_diff = graph.call_function(
-                    torch.ops.aten.pow.Tensor_Scalar,
-                    (layer_deviation, 2),
-                )
-                var = graph.call_function(
-                    torch.ops.aten.mean.dim,
-                    (layer_sqr_diff, [-1]),
-                )
-                var_eps = graph.call_function(
-                    torch.ops.aten.add.Tensor,
-                    (var, eps),
-                )
-                rstd = graph.call_function(
-                    torch.ops.aten.rsqrt.default,
-                    (var_eps,),
-                )
-                rstd_reshape = graph.call_function(
-                    torch.ops.aten.view.default,
-                    (rstd, [layer_size, 1]),
-                )
-                layer_norm = graph.call_function(
-                    torch.ops.aten.mul.Tensor,
-                    (layer_deviation, rstd_reshape),
-                )
-                layer_norm = graph.call_function(
-                    torch.ops.aten.view.default,
-                    (layer_norm, x_shape),
-                )
+                # Branch only on whether a reshape is needed; the normalization is shared.
+                if norm_size != x_shape[-1]:
+                    # Pack groups so that the last dimension equals norm_size.
+                    packed = graph.call_function(
+                        torch.ops.aten.reshape.default, (x, pack_shape)
+                    )
+                    normed = self._insert_norm(graph, packed, eps)
+                    # Restore the original shape after normalization.
+                    layer_norm = graph.call_function(
+                        torch.ops.aten.reshape.default, (normed, x_shape)
+                    )
+                else:
+                    # The input already has norm_size in the last dimension.
+                    layer_norm = self._insert_norm(graph, x, eps)
 
                 # weight
                 if weight:

tico/serialize/operators/op_mean.py

@@ -45,7 +45,7 @@ class MeanVisitor(NodeVisitor):
         args = MeanDimArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
         input = args.input
         dim = args.dim
-        keep_dims = args.keep_dims
+        keep_dims = args.keepdim
 
         dim_i32 = circle_legalize_dtype_to(dim, dtype=torch.int32)
         inputs = [input, dim_i32]

tico/utils/validate_args_kwargs.py

@@ -599,7 +599,7 @@ class MeanDimArgs:
 
     input: torch.fx.Node
     dim: List[int]
-    keep_dims: bool = False
+    keepdim: bool = False
     dtype: Optional[torch.dtype] = None
 
 

{tico-0.1.0.dev250520.dist-info → tico-0.1.0.dev250521.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tico
-Version: 0.1.0.dev250520
+Version: 0.1.0.dev250521
 Summary: Convert exported Torch module to circle
 Home-page: UNKNOWN
 License: UNKNOWN

{tico-0.1.0.dev250520.dist-info → tico-0.1.0.dev250521.dist-info}/RECORD

@@ -1,4 +1,4 @@
-tico/__init__.py,sha256=80nv6fiJM89LQkz5EODIec-qHWe6hHtEOZ5u7L1wYnY,1181
+tico/__init__.py,sha256=zsOHhqeEvcmCBCvdS74zuPFmwmyF7Ls-Y-7EGShtpIg,1181
 tico/pt2_to_circle.py,sha256=PPmFNw20jw2Z2VyM3ln9pX__jTzBOAZiv0gT5a-p-Y8,2666
 tico/config/__init__.py,sha256=xZzCXjZ84qE-CsBi-dfaL05bqpQ3stKKfTXhnrJRyVs,142
 tico/config/base.py,sha256=anwOiJFkUxUi7Cef573JgQcjk6S-FSi6O_TLjYASW-g,1244
@@ -70,7 +70,7 @@ tico/passes/decompose_addmm.py,sha256=_yNX7wx1Y9HJI5ksUJI-UQLHpoNawbUbF8kcm2zGHw
 tico/passes/decompose_batch_norm.py,sha256=d1V9UOkm_5BV0NGLyuQfz4I9NpO7I3ZrRugt7EXM-XM,7016
 tico/passes/decompose_fake_quantize.py,sha256=7ZJyTIDj2iKgWa5q8mBSq6k0GX0vs_XyQdsIiWFJoTU,5175
 tico/passes/decompose_fake_quantize_tensor_qparams.py,sha256=kOQaODKl_GCE19h-UZGmxnTcHtlvphI63dVAmMQL_Bk,13823
-tico/passes/decompose_group_norm.py,sha256=UVEzOBi4aUr3IlP6iWlYlo4u-8cmJ4JvKl8r2-qGBj4,9670
+tico/passes/decompose_group_norm.py,sha256=xn1xnT-2e6BvelRAzX8O7wg9kBWURmPldkRvpfYFXHQ,9407
 tico/passes/decompose_grouped_conv2d.py,sha256=KJhH6PX7l9k9T8KBV8JDAvaSfJuUnRo_jtvGF2aM-LA,8277
 tico/passes/decompose_slice_scatter.py,sha256=ko9p8v-zY5rOx4aSpWomwSdSWb1lIF32gnU7ik5xgII,5604
 tico/passes/extract_dtype_kwargs.py,sha256=hfGJ_GfZULbBmLif2AJkhPHVifhucxBiLoQI862Yejk,4303
@@ -140,7 +140,7 @@ tico/serialize/operators/op_logical_not.py,sha256=ugrVcRqR3IvUUaiRVW5cArCYJbzmkc
 tico/serialize/operators/op_lt.py,sha256=_vA7dWpV9wVBxB7JL9pLQT9BsV91NGQBq_0auAtHK5Y,2080
 tico/serialize/operators/op_max_pool2d_with_indices.py,sha256=ilQdirgSOjJR6dRIgAEF-sFmjPLQB3O2F_Fq5mbpYNA,5203
 tico/serialize/operators/op_maximum.py,sha256=JjBr6gWEnuakLuk1_feotTHfIIm3s5YqWmqhUMpSPI0,1873
-tico/serialize/operators/op_mean.py,sha256=e7uRPXYHrq3lH7vANcjCHXTRdMWZgiGFBmFdGaOQmts,2269
+tico/serialize/operators/op_mean.py,sha256=rVQZOxCJkHFY4kQBAS1HVK0HkcqxgkSy6zvEDLX_WYQ,2267
 tico/serialize/operators/op_minimum.py,sha256=fASjQVcTPCin02umQwFPdq2ss-Ve7S5A33J3QmmQ_wQ,1873
 tico/serialize/operators/op_mm.py,sha256=fHggR9dmlwXw0DAyn__2JbG7e0q1Jhfmi5-2jDlpRDk,6730
 tico/serialize/operators/op_mul.py,sha256=42Guc0MWBGBCZoj9-4LcLtTMtUPwsmDSVmvkR8tqLhM,3165
@@ -187,14 +187,14 @@ tico/utils/passes.py,sha256=kGmDe__5cPaO6i5EDAoXSVe6yXEoX9hAny4ROb3ZEmQ,2409
 tico/utils/register_custom_op.py,sha256=iRQvdqlBqrJxq_pNkvJyDIJD_SYtCUl88wwbbuvSwlk,22952
 tico/utils/trace_decorators.py,sha256=ddLIiKQfSaQrxgF1kNpwjFTQnXENzeSfcr1kuAW4jGI,3221
 tico/utils/utils.py,sha256=pybDU1LoNhjEplANig11lboX9yzYRkvFCSmyYth_2Do,10359
-tico/utils/validate_args_kwargs.py,sha256=krT68b5CfBI9rxBIOsgYSy0LfEJqLfKfRikkp8ep9oQ,24726
+tico/utils/validate_args_kwargs.py,sha256=z7JySDhyviI7G6Mo54AFdTiJIOE6Q6UrYnOFJj-5M24,24724
 tico/utils/mx/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/utils/mx/elemwise_ops.py,sha256=V6glyAHsVR1joqpsgnNytatCD_ew92xNWZ19UFDoMTA,10281
 tico/utils/mx/formats.py,sha256=uzNWyu-1onUlwQfX5cZ6fZSUfHMRqorper7_T1k3jfk,3404
 tico/utils/mx/mx_ops.py,sha256=RcfUTYVi-wilGB2sC35OeARdwDqnixv7dG5iyZ-fQT8,8555
-tico-0.1.0.dev250520.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
-tico-0.1.0.dev250520.dist-info/METADATA,sha256=E3Tizp7XTBYhstti5Qh5Q5HYd-7XiVA6zfwUD8T0DII,8633
-tico-0.1.0.dev250520.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-tico-0.1.0.dev250520.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
-tico-0.1.0.dev250520.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
-tico-0.1.0.dev250520.dist-info/RECORD,,
+tico-0.1.0.dev250521.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
+tico-0.1.0.dev250521.dist-info/METADATA,sha256=3DnR9WqN27XltH1Abj77Oc2hEgKbT99QZ0w0oVdvAuc,8633
+tico-0.1.0.dev250521.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+tico-0.1.0.dev250521.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
+tico-0.1.0.dev250521.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
+tico-0.1.0.dev250521.dist-info/RECORD,,