support shift by Tensor (#15623)

* support shift by Tensor

* use mixin

extra/torch_backend/backend.py

@@ -511,9 +511,8 @@ tiny_backend_out = {**{f"aten.{x}.out":getattr(Tensor,x) for x in simple_tensor_
   "aten.fmod.Tensor_out": lambda input,other: input-input.div(other, rounding_mode="trunc")*other,
   # TODO: this might result in overflow issues
   "aten.round.decimals_out": lambda self,decimals: (self*10**decimals).round()/10**decimals,
-  # TODO: support this in tinygrad. shift by Tensor not supported
-  "aten.bitwise_left_shift.Tensor_out": lambda x,y: x*(2**y),
-  "aten.bitwise_right_shift.Tensor_out": lambda x,y: x//(2**y),
+  "aten.bitwise_left_shift.Tensor_out": lambda x,y: x<<y,
+  "aten.bitwise_right_shift.Tensor_out": lambda x,y: x>>y,
   # not in tinygrad. are there decomps for these?
   "aten.log1p.out": lambda self: (self+1).log(),
   "aten.expm1.out": lambda self: self.exp() - 1,

tinygrad/mixin/elementwise.py

@@ -255,9 +255,25 @@ class ElementwiseMixin(DTypeMixin, CreationMixin):
   # NOTE: __eq__ isn't overridden, and means the same thing as is by default
 
   def lshift(self, x: Self | int, reverse: bool = False) -> Self:
+    """
+    Computes left arithmetic shift of `self` by `x` bits. `self` must have integer dtype.
+    Equivalent to `self << x`.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([1, 3, 31], dtype=dtypes.uint8).lshift(2).numpy())
+    ```
+    """
     return self._binop(Ops.SHL, x, reverse)
 
   def rshift(self, x: Self | int, reverse: bool = False) -> Self:
+    """
+    Computes right arithmetic shift of `self` by `x` bits. `self` must have integer dtype.
+    Equivalent to `self >> x`.
+
+    ```python exec="true" source="above" session="tensor" result="python"
+    print(Tensor([4, 13, 125], dtype=dtypes.uint8).rshift(2).numpy())
+    ```
+    """
     return self._binop(Ops.SHR, x, reverse)
 
   def __lshift__(self, x: Self | int) -> Self:

tinygrad/tensor.py

@@ -2511,30 +2511,6 @@ class Tensor(OpMixin):
     a, b = self._broadcasted(x, reverse)
     return a - a.div(b, rounding_mode="floor") * b
 
-  def lshift(self, x:Tensor|int, reverse=False) -> Tensor:
-    """
-    Computes left arithmetic shift of `self` by `x` bits. `self` must have integer dtype.
-    Equivalent to `self << x`.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor([1, 3, 31], dtype=dtypes.uint8).lshift(2).numpy())
-    ```
-    """
-    assert dtypes.is_int(self.dtype) and isinstance(x, int) and x >= 0 and not reverse, f"not supported {self.dtype=} {x=}"
-    return self.mul(2 ** x, reverse)
-
-  def rshift(self, x:Tensor|int, reverse=False) -> Tensor:
-    """
-    Computes right arithmetic shift of `self` by `x` bits. `self` must have integer dtype.
-    Equivalent to `self >> x`.
-
-    ```python exec="true" source="above" session="tensor" result="python"
-    print(Tensor([4, 13, 125], dtype=dtypes.uint8).rshift(2).numpy())
-    ```
-    """
-    assert dtypes.is_int(self.dtype) and isinstance(x, int) and x >= 0 and not reverse, f"not supported {self.dtype=} {x=}"
-    return (self & ~(2**x - 1)).idiv(2 ** x)
-
   def where(self:Tensor, x:Tensor|ConstType|sint, y:Tensor|ConstType|sint) -> Tensor:
     """
     Returns a tensor of elements selected from either `x` or `y`, depending on `self`.