From a05187788d2db7f4c97256c2eec8cb6ff294dfb7 Mon Sep 17 00:00:00 2001 From: Om Singhal Date: Mon, 7 Sep 2026 17:23:20 -0400 Subject: [PATCH] Fix aten_pow_scalar type promotion when the exponent is not floating point aten_pow_scalar cast the scalar base down to the exponent's dtype, so a float base over an integer or boolean exponent built an integer Pow. torch promotes the other way: a float scalar outranks an integral tensor, so 2.0 ** torch.tensor([1, 2, 3]) is float32. The exporter stamps the float result type on the output value while the node itself produces int64, so the model fails to load in onnxruntime. Boolean exponents fail earlier still, since Pow has no boolean inputs. Promote to float32 when a float scalar meets an integral exponent, and to int64 when an int scalar meets a boolean one. Every case that already agreed with torch keeps the exponent's dtype and the same nodes as before. --- .../function_libs/torch_lib/ops/core.py | 10 ++++ .../function_libs/torch_lib/e2e_ops_tests.py | 59 +++++++++++++++++++ 2 files changed, 69 insertions(+) diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py index 27d30e377f..22abeee6d2 100644 --- a/onnxscript/function_libs/torch_lib/ops/core.py +++ b/onnxscript/function_libs/torch_lib/ops/core.py @@ -8029,6 +8029,16 @@ def aten_pow_tensor_scalar(self: TReal, exponent: float) -> TReal: @torch_op("aten::pow.Scalar", trace_only=True) def aten_pow_scalar(self: float, exponent: TTensor) -> TTensor: """pow.Scalar(Scalar self, Tensor exponent) -> Tensor""" + if not isinstance(self, int) and not exponent.dtype.is_floating_point(): + # A float scalar outranks an integral exponent, so torch promotes the result to + # the default float type instead of narrowing the scalar down to the exponent + return op.Pow(op.Cast(self, to=FLOAT.dtype), op.Cast(exponent, to=FLOAT.dtype)) + if exponent.dtype == ir.DataType.BOOL: + # Pow has no boolean inputs, and an int scalar over a boolean exponent + # promotes to the default int type in torch + return op.Pow(op.Cast(self, to=INT64.dtype), op.Cast(exponent, to=INT64.dtype)) + # The exponent is in the same or a higher type category than the scalar, so it + # decides the result type. e.g. 2.0 ** float16 tensor is float16 return op.Pow(op.Cast(self, to=exponent.dtype), exponent) diff --git a/tests/function_libs/torch_lib/e2e_ops_tests.py b/tests/function_libs/torch_lib/e2e_ops_tests.py index 8a50b5d58d..bfbe33e434 100644 --- a/tests/function_libs/torch_lib/e2e_ops_tests.py +++ b/tests/function_libs/torch_lib/e2e_ops_tests.py @@ -84,6 +84,65 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: ) _testing.assert_onnx_program(onnx_program) + def test_pow_scalar_float_int(self): + class PowModel(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return 2.0**x + + onnx_program = torch.onnx.export( + PowModel(), + (torch.tensor([1, 2, 3], dtype=torch.int64),), + dynamo=True, + optimize=False, + ) + _testing.assert_onnx_program(onnx_program) + + def test_pow_scalar_float_bool(self): + class PowModel(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return 2.0**x + + onnx_program = torch.onnx.export( + PowModel(), (torch.tensor([True, False]),), dynamo=True, optimize=False + ) + _testing.assert_onnx_program(onnx_program) + + def test_pow_scalar_float_float16(self): + class PowModel(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return 2.0**x + + onnx_program = torch.onnx.export( + PowModel(), + (torch.tensor([1.0, 2.0], dtype=torch.float16),), + dynamo=True, + optimize=False, + ) + _testing.assert_onnx_program(onnx_program) + + def test_pow_scalar_int_int(self): + class PowModel(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return 2**x + + onnx_program = torch.onnx.export( + PowModel(), + (torch.tensor([1, 2, 3], dtype=torch.int64),), + dynamo=True, + optimize=False, + ) + _testing.assert_onnx_program(onnx_program) + + def test_pow_scalar_int_bool(self): + class PowModel(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return 2**x + + onnx_program = torch.onnx.export( + PowModel(), (torch.tensor([True, False]),), dynamo=True, optimize=False + ) + _testing.assert_onnx_program(onnx_program) + def test_mul_tensor_scalar_float(self): class Model(torch.nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: