[Muhammad Usman Mateen (usmanmateen)]

x << s and x >> s with a tensor s fail to export:

DispatchError: No ONNX function found for <OpOverload(op='aten.__lshift__', overload='Tensor')>

torch.bitwise_left_shift(x, s) and torch.bitwise_right_shift(x, s) export fine, as do x << 2 and x >> 2 (the .Scalar overloads, #1645). Only the operators with a tensor shift amount are missing. A shift tensor is typical when unpacking packed low-bit weights, e.g. (packed[:, :, None] >> shifts) & 0xF.

This registers aten::__lshift__.Tensor and aten::__rshift__.Tensor on the existing aten_bitwise_left_shift and aten_bitwise_right_shift.

Tests

• test_shift_operators_with_tensor_shift (int8, int32 and int64, with negative values) and test_unpack_int4_with_shift_operator in e2e_ops_tests.py. Both fail on main with the error above and pass with the change.
• e2e_ops_tests.py -k "shift or bitwise": 10 passed. ops_test.py -k "shift or bitwise": 16 passed, 6 skipped.
• ruff check and ruff format --check (0.15.8) pass.

I also exported the operators and the torch.bitwise_*_shift functions with the same inputs: int8, int16, int32, int64 and uint8, shift amounts at and above the bit width, 0-d and broadcast shift tensors, and empty inputs. In every case the operators now give the same result as the functions, or the same error. Two things fail the same way for both forms, before and after this change:

• 16-bit types: ONNX Runtime has no BitShift kernel for them.
• uint8: the result comes back as int8 (Keep unsigned dtypes in bitwise shifts #3088).

Not covered here

Operands of different dtypes still fail for the operators, e.g. packed_int32 >> torch.arange(0, 32, 4) (int64 shifts). PyTorch's exporter type-promotion table (torch/onnx/_internal/fx/passes/type_promotion.py) has rules for bitwise_left_shift and bitwise_right_shift but not for __lshift__ and __rshift__, so torchlib gets the operands unpromoted and the dtype assertion fails. With this change, adding those two rules with ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT (the same as for bitwise_left_shift) made the mixed cases I tried match eager, including 0-d shifts and the arange unpacking above.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 73.07%. Comparing base (06f8fc5) to head (344ea09).
⚠️ Report is 2 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #3109   +/-   ##
=======================================
  Coverage   73.07%   73.07%           
=======================================
  Files         267      267           
  Lines       32722    32722           
  Branches     3114     3114           
=======================================
  Hits        23913    23913           
  Misses       7767     7767           
  Partials     1042     1042           

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.