Keep synthesized layer_norm gamma/beta at normalized_shape

coreai_torch/_aten_to_core.py

@@ -1970,11 +1970,11 @@ def replace_layer_norm(
     bias = None if node.args[3] is None else _get_operand(values_map, node, 3)
     eps = node.args[4]
 
-    numel = int(np.prod(normalized_shape))
+    # Identity gamma/beta keep normalized_shape to match the declared `axes`.
     if weight is None:
-        weight = coreai.constant([1.0] * numel, dtype=np.float32)
+        weight = coreai.constant(np.ones(normalized_shape, dtype=np.float32))
     if bias is None:
-        bias = coreai.constant([0.0] * numel, dtype=np.float32)
+        bias = coreai.constant(np.zeros(normalized_shape, dtype=np.float32))
 
     input_rank = x.type.rank
     input_ele_type = x.type.element_type
@@ -2022,11 +2022,7 @@ def layer_norm(input: Value, gamma: Value, beta: Value) -> Value:
         weight = coreai.cast(gamma, input_ele_type)
         bias = coreai.cast(beta, input_ele_type)
 
-        # Reshape for multi-dim normalized_shape (e.g. [32] → [4, 8]).
-        if len(normalized_shape) > 1:
-            weight = coreai.reshape(weight, list(normalized_shape))
-            bias = coreai.reshape(bias, list(normalized_shape))
-
+        # gamma/beta always arrive shaped like normalized_shape.
         # Broadcast gamma/beta to match the norm output shape.
         norm_shape = coreai.get_shape(norm)
         w_shape = coreai.constant(list(normalized_shape), dtype=np.uint32)

tests/ops/test_ops_ir.py

@@ -6052,6 +6052,39 @@ def forward(self, x: Tensor) -> Tensor:
             """,
         )
 
+    def test_multi_dim_normalized_shape_without_affine(self) -> None:
+        """Synthesized gamma/beta keep normalized_shape, matching `axes`."""
+
+        class LayerNormModel(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.ln = nn.LayerNorm((4, 8), elementwise_affine=False)
+
+            def forward(self, x: Tensor) -> Tensor:
+                return self.ln(x)
+
+        ir = get_ir(LayerNormModel().eval(), x=torch.rand(2, 4, 8))
+        filecheck_pattern(
+            ir,
+            check_file="""
+                // CHECK-LABEL: module {
+                // CHECK-NEXT:   coreai.graph private noinline @layer_norm_{{.*}}(%[[INPUT:.*]]: tensor<2x4x8xf32> {coreai.name = "input"}, %[[GAMMA:.*]]: tensor<4x8xf32> {coreai.name = "gamma"}, %[[BETA:.*]]: tensor<4x8xf32> {coreai.name = "beta"}) -> tensor<2x4x8xf32> attributes {__coreai_pure__, composite_decl = #coreai.composite_declaration<"layer_norm" = {input_names = ["input", "gamma", "beta"], op_attrs = {axes = [1 : si64, 2 : si64], eps = 9.99999974E-6 : f32, version = 1 : si64}, output_names = ["output"]}>, template_op = "layer_norm"} {
+                // CHECK-NOT:       coreai.reshape
+                // CHECK:           %[[NORM:.*]] = coreai.decomposable.broadcasting_mul %{{.*}}, %{{.*}} : (tensor<2x4x8xf32>, tensor<2x1x1xf32>) -> tensor<2x4x8xf32>
+                // CHECK-NEXT:      %[[SCALED:.*]] = coreai.decomposable.broadcasting_mul %[[NORM]], %[[GAMMA]] : (tensor<2x4x8xf32>, tensor<4x8xf32>) -> tensor<2x4x8xf32>
+                // CHECK-NEXT:      %[[SHIFTED:.*]] = coreai.decomposable.broadcasting_add %[[SCALED]], %[[BETA]] : (tensor<2x4x8xf32>, tensor<4x8xf32>) -> tensor<2x4x8xf32>
+                // CHECK-NEXT:      coreai.output %[[SHIFTED]] : tensor<2x4x8xf32>
+                // CHECK-NEXT:    }
+                // CHECK-NEXT:    coreai.graph @main(%[[X:.*]]: tensor<2x4x8xf32> {coreai.name = "x"}) -> (tensor<2x4x8xf32> {coreai.name = "{{.*}}"}) attributes {__coreai_pure__} {
+                // CHECK-NEXT:      %[[ONE:.*]] = coreai.constant dense<1.000000e+00> : tensor<4x8xf32>
+                // CHECK-NEXT:      %[[ZERO:.*]] = coreai.constant dense<0.000000e+00> : tensor<4x8xf32>
+                // CHECK-NEXT:      %[[R:.*]] = coreai.invoke @layer_norm_{{.*}}(%[[X]], %[[ONE]], %[[ZERO]])  : (tensor<2x4x8xf32>, tensor<4x8xf32>, tensor<4x8xf32>) -> tensor<2x4x8xf32>
+                // CHECK-NEXT:      coreai.output %[[R]] : tensor<2x4x8xf32>
+                // CHECK-NEXT:    }
+                // CHECK-NEXT:  }
+            """,
+        )
+
 
 class TestNeScalarIR:
     def test_static(self) -> None: