From e13d98a941f8f4e76fb9f1d9d8cce714562d98e5 Mon Sep 17 00:00:00 2001 From: Devin Lai Date: Thu, 3 Sep 2026 07:56:51 -0400 Subject: [PATCH] linalg_vector_norm: match p-norm exponent dtype The general p-norm lowering passed Python scalars directly to broadcasting_pow, which could create exponent constants with a different element type from the operand. Integer ord values and float16 inputs therefore produced programs that failed verification. Normalize ord to a float and build both exponent constants in the operand element type. Add end-to-end coverage for integer ord spellings and float16 inputs. --- coreai_torch/_aten_to_core.py | 20 ++++++++++++++++---- tests/ops/test_ops.py | 11 ++++++++++- 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/coreai_torch/_aten_to_core.py b/coreai_torch/_aten_to_core.py index 1af18a3..a45c76d 100644 --- a/coreai_torch/_aten_to_core.py +++ b/coreai_torch/_aten_to_core.py @@ -2078,7 +2078,11 @@ def replace_linalg_vector_norm( """Linalg vector norm: L0/L1/L2/Linf/Lp norms along specified dims.""" x = _get_operand(values_map, node, 0) args = node.args - ord_val = args[1] if len(args) > 1 and args[1] is not None else 2.0 + # `ord` is a Scalar in the ATen schema, so an integer spelling such as + # `ord=3` arrives as a Python int. Normalize it: the general p-norm branch + # below feeds it to broadcasting_pow, where an int would build an integer + # exponent constant that does not match the float operand's element type. + ord_val = float(args[1]) if len(args) > 1 and args[1] is not None else 2.0 dim = args[2] if len(args) > 2 and args[2] is not None else None keepdim = args[3] if len(args) > 3 and args[3] is not None else False @@ -2118,12 +2122,20 @@ def linalg_vector_norm(input: Value) -> Value: elif ord_val == float("-inf"): result = coreai.reduce_min(coreai.abs_(input), dims) else: - # General p-norm: (sum(|x|^p))^(1/p) + # General p-norm: (sum(|x|^p))^(1/p). broadcasting_pow requires the + # exponent's element type to match the operand's, so build both + # exponents as constants of the input's element type rather than + # letting the Python scalar pick one. + element_type = input.type.element_type result = coreai.broadcasting_pow( coreai.reduce_sum( - coreai.broadcasting_pow(coreai.abs_(input), ord_val), dims + coreai.broadcasting_pow( + coreai.abs_(input), + coreai.constant(ord_val, dtype=element_type), + ), + dims, ), - 1.0 / ord_val, + coreai.constant(1.0 / ord_val, dtype=element_type), ) return result if keepdim else coreai.shrink_dims(result, dims) diff --git a/tests/ops/test_ops.py b/tests/ops/test_ops.py index d8f8233..4681b88 100644 --- a/tests/ops/test_ops.py +++ b/tests/ops/test_ops.py @@ -2743,6 +2743,15 @@ def forward(self, x: Tensor) -> Tensor: # Large order (treated as infinity norm) (torch.rand(5), None, 15.0, False), (torch.rand(2, 3, 4), 2, -3.0, False), + # Integer `ord`: the ATen schema types it as a Scalar, so `ord=3` stays + # an int and the p-norm exponent constant must not be built from it + # verbatim. + (torch.rand(2, 3, 4) + 0.5, 2, 3, False), + (torch.rand(3, 4) + 0.5, 1, -1, True), + (torch.rand(5) + 0.5, 0, -2, False), + # fp16 operand through the general p-norm branch + (torch.rand(3, 4, dtype=torch.float16) + 0.5, 1, 3.0, False), + (torch.rand(2, 3, 4, dtype=torch.float16) + 0.5, 2, 3, True), ], ) @pytest.mark.parametrize( @@ -2751,7 +2760,7 @@ def forward(self, x: Tensor) -> Tensor: async def test_linalg_vector_norm( x: Tensor, dim: int | list[int] | None, - ord: float, + ord: int | float, keepdim: bool, dynamic_dims: tuple[int], ) -> None: