Match vector norm exponent constants to the input dtype

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)

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: