Keep unsigned dtypes in bitwise shifts

[Muhammad Usman Mateen (usmanmateen)]

aten_bitwise_left_shift and aten_bitwise_right_shift cast the inputs to the unsigned type of the same width, shift, and then always cast the result to the signed type. For signed inputs that is the intent, but for uint8 the result comes back as int8:

import torch


class Model(torch.nn.Module):
    def forward(self, x):
        return x << 4, x >> 4


x = torch.tensor([0x0F, 0x21, 0xF0, 0xFF], dtype=torch.uint8)
onnx_program = torch.onnx.export(Model(), (x,), dynamo=True)
print(Model()(x))
# (tensor([240,  16,   0, 240], dtype=torch.uint8), tensor([ 0,  2, 15, 15], dtype=torch.uint8))
print(onnx_program(x))
# (tensor([-16,  16,   0, -16], dtype=torch.int8), tensor([ 0,  2, 15, 15], dtype=torch.int8))

Packing two 4-bit values into a byte, (hi << 4) | lo with uint8 tensors, exports a model that ONNX Runtime does not load:

Type Error: Type parameter (T) of Optype (BitwiseOr) bound to different types (tensor(int8) and tensor(uint8) in node (node_bitwise_or).

ONNX BitShift supports the unsigned types directly, and a logical shift is what PyTorch does for them, so this returns BitShift on the inputs as they are when the dtype is unsigned.

Tests

test_bitwise_shift_uint8 in tests/function_libs/torch_lib/e2e_ops_tests.py exports x << 4, x >> 4, bitwise_left_shift and bitwise_right_shift with a tensor shift amount, and (x << 4) | y, all on uint8, and checks them against eager. It fails on main with the error above and passes with this change.

• python -m pytest tests/function_libs/torch_lib/e2e_ops_tests.py -k bitwise: 2 passed
• python -m pytest tests/function_libs/torch_lib/ops_test.py -k bitwise: 16 passed, 6 skipped (126 subtests passed)
• ruff check and ruff format --check (0.15.8) on the changed files

Tested with torch 2.14.1+cpu and onnxruntime 1.31.0.