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: