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: