[torchlib] Register aten::__lshift__ and __rshift__

justinchuby · justinchuby · commit 9020774589e4 · 2025-03-13T09:17:23.000-07:00
diff --git a/onnxscript/function_libs/torch_lib/ops/core.py b/onnxscript/function_libs/torch_lib/ops/core.py
@@ -1243,6 +1243,7 @@ def aten_bitwise_and(self: TInt, other: TInt) -> TInt:
         "aten::bitwise_left_shift.Tensor_Scalar",
         "aten::bitwise_left_shift.Scalar_Tensor",
         "_operator::__lshift__",
+        "aten::__lshift__.Scalar",
     ),
     trace_only=True,
 )
@@ -1263,6 +1264,7 @@ def aten_bitwise_left_shift_int16(self: INT16, other: INT16) -> INT16:
         "aten::bitwise_left_shift.Tensor_Scalar",
         "aten::bitwise_left_shift.Scalar_Tensor",
         "_operator::__lshift__",
+        "aten::__lshift__.Scalar",
     ),
     trace_only=True,
 )
@@ -1283,6 +1285,7 @@ def aten_bitwise_left_shift_int32(self: INT32, other: INT32) -> INT32:
         "aten::bitwise_left_shift.Tensor_Scalar",
         "aten::bitwise_left_shift.Scalar_Tensor",
         "_operator::__lshift__",
+        "aten::__lshift__.Scalar",
     ),
     trace_only=True,
 )
@@ -1303,6 +1306,7 @@ def aten_bitwise_left_shift_int64(self: INT64, other: INT64) -> INT64:
         "aten::bitwise_left_shift.Tensor_Scalar",
         "aten::bitwise_left_shift.Scalar_Tensor",
         "_operator::__lshift__",
+        "aten::__lshift__.Scalar",
     ),
     trace_only=True,
 )
@@ -1347,6 +1351,7 @@ def aten_bitwise_or(self: TInt, other: TInt) -> TInt:
         "aten::bitwise_right_shift.Tensor_Scalar",
         "aten::bitwise_right_shift.Scalar_Tensor",
         "_operator::__rshift__",
+        "aten::__rshift__.Scalar",
     )
 )
 def aten_bitwise_right_shift_int16(self: INT16, other: INT16) -> INT16:
@@ -1377,6 +1382,7 @@ def aten_bitwise_right_shift_int16(self: INT16, other: INT16) -> INT16:
         "aten::bitwise_right_shift.Tensor_Scalar",
         "aten::bitwise_right_shift.Scalar_Tensor",
         "_operator::__rshift__",
+        "aten::__rshift__.Scalar",
     )
 )
 def aten_bitwise_right_shift_int32(self: INT32, other: INT32) -> INT32:
@@ -1407,6 +1413,7 @@ def aten_bitwise_right_shift_int32(self: INT32, other: INT32) -> INT32:
         "aten::bitwise_right_shift.Tensor_Scalar",
         "aten::bitwise_right_shift.Scalar_Tensor",
         "_operator::__rshift__",
+        "aten::__rshift__.Scalar",
     )
 )
 def aten_bitwise_right_shift_int64(self: INT64, other: INT64) -> INT64:
@@ -1440,6 +1447,7 @@ def aten_bitwise_right_shift_int64(self: INT64, other: INT64) -> INT64:
         "aten::bitwise_right_shift.Tensor_Scalar",
         "aten::bitwise_right_shift.Scalar_Tensor",
         "_operator::__rshift__",
+        "aten::__rshift__.Scalar",
     )
 )
 def aten_bitwise_right_shift_int8(self: INT8, other: INT8) -> INT8:

Original file line number	Diff line number	Diff line change
`@@ -1243,6 +1243,7 @@ def aten_bitwise_and(self: TInt, other: TInt) -> TInt:`
`1243`	`1243`	`"aten::bitwise_left_shift.Tensor_Scalar",`
`1244`	`1244`	`"aten::bitwise_left_shift.Scalar_Tensor",`
`1245`	`1245`	`"_operator::__lshift__",`
	`1246`	`+ "aten::__lshift__.Scalar",`
`1246`	`1247`	`),`
`1247`	`1248`	`trace_only=True,`
`1248`	`1249`	`)`
`@@ -1263,6 +1264,7 @@ def aten_bitwise_left_shift_int16(self: INT16, other: INT16) -> INT16:`
`1263`	`1264`	`"aten::bitwise_left_shift.Tensor_Scalar",`
`1264`	`1265`	`"aten::bitwise_left_shift.Scalar_Tensor",`
`1265`	`1266`	`"_operator::__lshift__",`
	`1267`	`+ "aten::__lshift__.Scalar",`
`1266`	`1268`	`),`
`1267`	`1269`	`trace_only=True,`
`1268`	`1270`	`)`
`@@ -1283,6 +1285,7 @@ def aten_bitwise_left_shift_int32(self: INT32, other: INT32) -> INT32:`
`1283`	`1285`	`"aten::bitwise_left_shift.Tensor_Scalar",`
`1284`	`1286`	`"aten::bitwise_left_shift.Scalar_Tensor",`
`1285`	`1287`	`"_operator::__lshift__",`
	`1288`	`+ "aten::__lshift__.Scalar",`
`1286`	`1289`	`),`
`1287`	`1290`	`trace_only=True,`
`1288`	`1291`	`)`
`@@ -1303,6 +1306,7 @@ def aten_bitwise_left_shift_int64(self: INT64, other: INT64) -> INT64:`
`1303`	`1306`	`"aten::bitwise_left_shift.Tensor_Scalar",`
`1304`	`1307`	`"aten::bitwise_left_shift.Scalar_Tensor",`
`1305`	`1308`	`"_operator::__lshift__",`
	`1309`	`+ "aten::__lshift__.Scalar",`
`1306`	`1310`	`),`
`1307`	`1311`	`trace_only=True,`
`1308`	`1312`	`)`
`@@ -1347,6 +1351,7 @@ def aten_bitwise_or(self: TInt, other: TInt) -> TInt:`
`1347`	`1351`	`"aten::bitwise_right_shift.Tensor_Scalar",`
`1348`	`1352`	`"aten::bitwise_right_shift.Scalar_Tensor",`
`1349`	`1353`	`"_operator::__rshift__",`
	`1354`	`+ "aten::__rshift__.Scalar",`
`1350`	`1355`	`)`
`1351`	`1356`	`)`
`1352`	`1357`	`def aten_bitwise_right_shift_int16(self: INT16, other: INT16) -> INT16:`
`@@ -1377,6 +1382,7 @@ def aten_bitwise_right_shift_int16(self: INT16, other: INT16) -> INT16:`
`1377`	`1382`	`"aten::bitwise_right_shift.Tensor_Scalar",`
`1378`	`1383`	`"aten::bitwise_right_shift.Scalar_Tensor",`
`1379`	`1384`	`"_operator::__rshift__",`
	`1385`	`+ "aten::__rshift__.Scalar",`
`1380`	`1386`	`)`
`1381`	`1387`	`)`
`1382`	`1388`	`def aten_bitwise_right_shift_int32(self: INT32, other: INT32) -> INT32:`
`@@ -1407,6 +1413,7 @@ def aten_bitwise_right_shift_int32(self: INT32, other: INT32) -> INT32:`
`1407`	`1413`	`"aten::bitwise_right_shift.Tensor_Scalar",`
`1408`	`1414`	`"aten::bitwise_right_shift.Scalar_Tensor",`
`1409`	`1415`	`"_operator::__rshift__",`
	`1416`	`+ "aten::__rshift__.Scalar",`
`1410`	`1417`	`)`
`1411`	`1418`	`)`
`1412`	`1419`	`def aten_bitwise_right_shift_int64(self: INT64, other: INT64) -> INT64:`
`@@ -1440,6 +1447,7 @@ def aten_bitwise_right_shift_int64(self: INT64, other: INT64) -> INT64:`
`1440`	`1447`	`"aten::bitwise_right_shift.Tensor_Scalar",`
`1441`	`1448`	`"aten::bitwise_right_shift.Scalar_Tensor",`
`1442`	`1449`	`"_operator::__rshift__",`
	`1450`	`+ "aten::__rshift__.Scalar",`
`1443`	`1451`	`)`
`1444`	`1452`	`)`
`1445`	`1453`	`def aten_bitwise_right_shift_int8(self: INT8, other: INT8) -> INT8:`