From 5ef1bd1a8e60abfd6cdbbebf956cdd7f6ff9cb3a Mon Sep 17 00:00:00 2001 From: gokulkrishna98 Date: Tue, 11 Aug 2026 09:55:53 -0700 Subject: [PATCH] Keep synthesized layer_norm gamma/beta at normalized_shape When elementwise_affine=False, replace_layer_norm built the identity gamma/beta as a flat numel-element vector, so a layer_norm over normalized_shape (4, 8) produced a composite taking tensor<32xf32> gamma/beta while its declaration said axes = [1, 2]. The body reshaped them back before broadcasting, so numerics were correct, but the composite boundary was inconsistent with the declared axes and differed from the elementwise_affine=True case, where ATen's real params arrive already shaped (4, 8). A consumer that implements the composite from its declaration instead of inlining the body sees a rank-1 gamma against two normalized axes. Build the constants with np.ones/np.zeros at normalized_shape so both paths present the same interface, and drop the now-dead reshape. --- coreai_torch/_aten_to_core.py | 12 ++++-------- tests/ops/test_ops_ir.py | 33 +++++++++++++++++++++++++++++++++ 2 files changed, 37 insertions(+), 8 deletions(-) diff --git a/coreai_torch/_aten_to_core.py b/coreai_torch/_aten_to_core.py index f8fe446..f129a82 100644 --- a/coreai_torch/_aten_to_core.py +++ b/coreai_torch/_aten_to_core.py @@ -1956,11 +1956,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 @@ -2008,11 +2008,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) diff --git a/tests/ops/test_ops_ir.py b/tests/ops/test_ops_ir.py index 1e242cb..ec2b737 100644 --- a/tests/ops/test_ops_ir.py +++ b/tests/ops/test_ops_ir.py @@ -6015,6 +6015,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: