Skip to content

Fix aten_pow_scalar type promotion when the exponent is not floating point - #3036

Open
om singhal (Om-singhaI) wants to merge 1 commit into
microsoft:mainfrom
Om-singhaI:fix/pow-scalar-type-promotion
Open

Fix aten_pow_scalar type promotion when the exponent is not floating point#3036
om singhal (Om-singhaI) wants to merge 1 commit into
microsoft:mainfrom
Om-singhaI:fix/pow-scalar-type-promotion

Conversation

@Om-singhaI

Copy link
Copy Markdown
Contributor

aten_pow_scalar casts the scalar base down to the exponent's dtype:

return op.Pow(op.Cast(self, to=exponent.dtype), exponent)

When the base is a Python float and the exponent tensor is an integer, that casts the float down to int64 and builds 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.

Repro on main:

import torch

class M(torch.nn.Module):
    def forward(self, x):
        return 2.0**x

x = torch.tensor([1, 2, 3])
onnx_program = torch.onnx.export(M(), (x,), dynamo=True, optimize=False)
onnx_program(x)
onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL :
Type Error: Type (tensor(float)) of output arg (pow_1) of node (node_pow_1)
does not match expected type (tensor(int64)).

The graph contradicts itself. The output value carries FLOAT because that's the dtype torch reports for the result, while the Cast we emit is to=7 and the Pow really returns int64. So it isn't just a wrong dtype, the model doesn't load at all.

A boolean exponent fails even earlier, for a float base and an int base both, because Pow accepts no boolean inputs:

[ONNXRuntimeError] : 10 : INVALID_GRAPH : This is an invalid model.
Type Error: Type 'tensor(bool)' of input parameter (val_1) of operator (Pow)
in node (node_pow_1) is invalid.

aten_pow_tensor_scalar right above already handles the mirror case correctly. It refuses to narrow and casts up to FLOAT instead. This applies the same rule in the other direction.

I promote to float32 when a float scalar meets an integral exponent, and to int64 when an int scalar meets a boolean one. Everything that already agreed with torch falls through to the original line untouched, so an int scalar over an int tensor still keeps the exponent's dtype and emits the same nodes as before.

Five tests in e2e_ops_tests.py next to the existing test_pow_tensor_scalar_* ones. Three of them fail on main (float base over int64, float base over bool, int base over bool). The other two pin the cases that must not move (float base over float16, int base over int64), and they pass either way.

…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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Development

Successfully merging this pull request may close these issues.

1 participant