Numeric overflow in JIT

[datumbox]

🐛 Bug

To Reproduce

The following snippet causes numeric overflow in JIT if x is uint8:

def invert(x):
    bound = 255.0
    return bound - x # Overflow happens here due to -x when uint8

x=torch.randint(0, 256, (10, 10), dtype=torch.uint8)

script_invert = torch.jit.script(invert)
invert(x).equal(script_invert(x)) # False

We can avoid it if we define the invert differently:

def invert2(x):
    bound = 255.0
    return -(x - bound) # Avoids underflow because x is casted to float before negated


script_invert2 = torch.jit.script(invert2)
invert2(x).equal(script_invert2(x)) # True

Expected behavior

The two implementations should yield equivalent results. The invert() should return the same result either it's scripted or not.

cc @gmagogsfm

[eellison]

Thanks for the issue. This should be fixed with #49247