From 43448e2989e4b22db971e3b8c4c02b9e93481458 Mon Sep 17 00:00:00 2001
From: sohailshk <122166523+sohailshk@users.noreply.github.com>
Date: Fri, 13 Nov 2026 01:51:32 +0530
Subject: [PATCH] [torchlib] Promote mixed dtypes in aten::pow.Tensor_Tensor

ONNX Pow returns the dtype of the base, so pow(int64 tensor, float tensor) was exported with an int64 result (values truncated) and pow(float32 tensor, float64 tensor) as float32, while torch promotes to the float dtype. Cast the operands to the dtype torch.result_type picks, treating zero-dim tensors as torch does.
---
 .../function_libs/torch_lib/ops/core.py       | 38 ++++++++++++++++++-
 .../function_libs/torch_lib/e2e_ops_tests.py  | 29 ++++++++++++++
 2 files changed, 66 insertions(+), 1 deletion(-)

diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py
index e7f1a611f7..87e4fc28f0 100644
--- a/onnxscript/function_libs/torch_lib/ops/core.py
+++ b/onnxscript/function_libs/torch_lib/ops/core.py
@@ -8062,10 +8062,46 @@ def aten_positive(self: TensorType) -> TensorType:
 @torch_op(("aten::pow.Tensor_Tensor", "_operator::pow"), trace_only=True)
 def aten_pow(self: TReal, exponent: TTensor) -> TReal:
     """pow(Tensor self, Tensor exponent) -> Tensor"""
-    # TODO(justinchuby): Add type promotion
+    if (
+        isinstance(self, ir.Value)
+        and isinstance(exponent, ir.Value)
+        and self.dtype != exponent.dtype
+        and (self.dtype.is_floating_point() or exponent.dtype.is_floating_point())
+    ):
+        # ONNX Pow keeps the dtype of the base, but torch promotes the result when
+        # the dtypes differ, e.g. an int64 base and a float32 exponent gives float32.
+        result_dtype = _promoted_float_dtype(self, exponent)
+        if self.dtype != result_dtype:
+            self = op.Cast(self, to=result_dtype)
+        # Pow accepts an integer exponent next to a float base, so integer exponents
+        # are left alone. Float and bool exponents (Pow has no bool input) follow the
+        # promoted dtype.
+        if not exponent.dtype.is_integer() and exponent.dtype != result_dtype:
+            exponent = op.Cast(exponent, to=result_dtype)
     return op.Pow(self, exponent)
 
 
+def _promoted_float_dtype(a: TensorType, b: TensorType) -> int:
+    """The result dtype torch gives to a binary op on tensors a and b, one of them float.
+
+    A zero-dim tensor only decides the dtype when it is of a higher category (float over
+    integer) than the other operand, mirroring torch.result_type.
+    """
+    a_float = a.dtype.is_floating_point()
+    b_float = b.dtype.is_floating_point()
+    if a_float != b_float:
+        # One float operand and one integer or bool operand: the float one decides.
+        return (a if a_float else b).dtype
+    a_scalar = a.shape is not None and len(a.shape) == 0
+    b_scalar = b.shape is not None and len(b.shape) == 0
+    if a_scalar != b_scalar:
+        # Both are floats of different widths. The dimensioned operand wins.
+        return (b if a_scalar else a).dtype
+    if {a.dtype, b.dtype} == {ir.DataType.FLOAT16, ir.DataType.BFLOAT16}:
+        return ir.DataType.FLOAT
+    return a.dtype if a.dtype.bitwidth >= b.dtype.bitwidth else b.dtype
+
+
 @torch_op("aten::pow.Tensor_Scalar", trace_only=True)
 def aten_pow_tensor_scalar(self: TReal, exponent: float) -> TReal:
     """pow(Tensor self, Scalar exponent) -> Tensor"""
diff --git a/tests/function_libs/torch_lib/e2e_ops_tests.py b/tests/function_libs/torch_lib/e2e_ops_tests.py
index 96cd89cf18..47b2040aca 100644
--- a/tests/function_libs/torch_lib/e2e_ops_tests.py
+++ b/tests/function_libs/torch_lib/e2e_ops_tests.py
@@ -45,6 +45,35 @@ def forward(self, x, indices, updates):
         onnx_program = torch.onnx.export(model, xs, dynamo=True)
         _testing.assert_onnx_program(onnx_program)
 
+    @parameterized.parameterized.expand(
+        [
+            ("int64_float32", torch.int64, torch.float32, (1, 3), (2, 1)),
+            ("float32_int64", torch.float32, torch.int64, (2, 1), (1, 3)),
+            ("float32_bool", torch.float32, torch.bool, (2, 1), (1, 3)),
+            ("float16_int64", torch.float16, torch.int64, (1, 3), (2, 1)),
+            ("float32_float64", torch.float32, torch.float64, (1, 3), (2, 1)),
+            ("float64_zero_dim_float32", torch.float64, torch.float32, (), (1, 3)),
+            ("int64_zero_dim_float64", torch.int64, torch.float64, (1, 3), ()),
+        ]
+    )
+    def test_pow_tensor_tensor_mixed_dtypes(
+        self, _: str, base_dtype, exponent_dtype, base_shape, exponent_shape
+    ):
+        class PowModel(torch.nn.Module):
+            def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+                return torch.pow(x, y)
+
+        def make(dtype, shape):
+            return (
+                torch.randint(1, 4, shape)
+                if not dtype.is_floating_point
+                else torch.rand(shape) + 0.5
+            ).to(dtype)
+
+        inputs = (make(base_dtype, base_shape), make(exponent_dtype, exponent_shape))
+        onnx_program = torch.onnx.export(PowModel(), inputs, dynamo=True, optimize=False)
+        _testing.assert_onnx_program(onnx_program)
+
     def test_pow_tensor_scalar_int_float(self):
         class PowModel(torch.nn.Module):
             def forward(self, x: torch.Tensor) -> torch.Tensor: