geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

move bitwise_and/bitwise_or/bitwise_xor to MathTrait [pr] (#10794)

Browse Source 
* move bitwise and, or, xor to MathTrait

* refactor
This commit is contained in:
Bhavya Gada
2025-06-17 09:19:43 -07:00
committed by GitHub
parent 531d143780
commit 3a474ef5b7
2 changed files with 47 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/tensor.py
-46
View File
@@ -3566,52 +3566,6 @@ class Tensor(MathTrait):
a, b = self._broadcasted(x, reverse)
return a - a.div(b, rounding_mode="floor") * b
def bitwise_xor(self, x:Tensor|ConstType, reverse=False) -> Tensor:
"""
Computes bitwise xor of `self` and `x`.
Equivalent to `self ^ x`.
Supports broadcasting to a common shape, type promotion, and integer, boolean inputs.
```python exec="true" source="above" session="tensor" result="python"
print(Tensor([-1, -2, 3]).bitwise_xor(Tensor([1, 0, 3])).numpy())
```
```python exec="true" source="above" session="tensor" result="python"
print(Tensor([True, True, False, False]).bitwise_xor(Tensor([True, False, True, False])).numpy())
```
"""
if self.dtype != dtypes.bool and not dtypes.is_int(self.dtype): raise RuntimeError(f"{self.dtype} is not supported")
return self._apply_broadcasted_uop(UOp.bitwise_xor, x, reverse)
def bitwise_and(self, x:Tensor|ConstType, reverse=False) -> Tensor:
"""
Computes the bitwise AND of `self` and `x`.
Equivalent to `self & x`.
Supports broadcasting to a common shape, type promotion, and integer, boolean inputs.
```python exec="true" source="above" session="tensor" result="python"
print(Tensor([2, 5, 255]).bitwise_and(Tensor([3, 14, 16])).numpy())
```
```python exec="true" source="above" session="tensor" result="python"
print(Tensor([True, True, False, False]).bitwise_and(Tensor([True, False, True, False])).numpy())
```
"""
if self.dtype != dtypes.bool and not dtypes.is_int(self.dtype): raise RuntimeError(f"{self.dtype} is not supported")
return self._apply_broadcasted_uop(UOp.bitwise_and, x, reverse)
def bitwise_or(self, x:Tensor|ConstType, reverse=False) -> Tensor:
"""
Computes the bitwise OR of `self` and `x`.
Equivalent to `self | x`.
Supports broadcasting to a common shape, type promotion, and integer, boolean inputs.
```python exec="true" source="above" session="tensor" result="python"
print(Tensor([2, 5, 255]).bitwise_or(Tensor([4, 4, 4])).numpy())
```
```python exec="true" source="above" session="tensor" result="python"
print(Tensor([True, True, False, False]).bitwise_or(Tensor([True, False, True, False])).numpy())
```
"""
if self.dtype != dtypes.bool and not dtypes.is_int(self.dtype): raise RuntimeError(f"{self.dtype} is not supported")
return self._apply_broadcasted_uop(UOp.bitwise_or, x, reverse)
def bitwise_not(self) -> Tensor:
"""
Computes the bitwise NOT of `self`.
tinygrad/uop/mathtraits.py
+47 -3
View File
@@ -14,6 +14,10 @@ class MathTrait:
def neg(self):
if (dtype:=getattr(self, 'dtype')) is None: raise TypeError(f"MathTraits __neg__ requires a dtype, {self=}")
return self.logical_not() if dtype.scalar() == dtypes.bool else self*(-1)
def _check_dtype(self):
if (dtype:=getattr(self, 'dtype')) is not None:
if isinstance(dtype, tuple): dtype = dtype[0]
if not (dtypes.is_bool(dtype) or dtypes.is_int(dtype)): raise RuntimeError(f"{dtype} is not supported")
def add(self, x, reverse=False):
"""
Adds `self` and `x`.
@@ -51,9 +55,49 @@ class MathTrait:
```
"""
return self._binop(Ops.MUL, x, reverse)
def bitwise_and(self, x, reverse=False): return self._binop(Ops.AND, x, reverse)
def bitwise_or(self, x, reverse=False): return self._binop(Ops.OR, x, reverse)
def bitwise_xor(self, x, reverse=False): return self._binop(Ops.XOR, x, reverse)
def bitwise_and(self, x, reverse=False):
"""
Computes the bitwise AND of `self` and `x`.
Equivalent to `self & x`.
Supports broadcasting to a common shape, type promotion, and integer, boolean inputs.
```python exec="true" source="above" session="tensor" result="python"
print(Tensor([2, 5, 255]).bitwise_and(Tensor([3, 14, 16])).numpy())
```
```python exec="true" source="above" session="tensor" result="python"
print(Tensor([True, True, False, False]).bitwise_and(Tensor([True, False, True, False])).numpy())
```
"""
self._check_dtype()
return self._binop(Ops.AND, x, reverse)
def bitwise_or(self, x, reverse=False):
"""
Computes the bitwise OR of `self` and `x`.
Equivalent to `self | x`.
Supports broadcasting to a common shape, type promotion, and integer, boolean inputs.
```python exec="true" source="above" session="tensor" result="python"
print(Tensor([2, 5, 255]).bitwise_or(Tensor([4, 4, 4])).numpy())
```
```python exec="true" source="above" session="tensor" result="python"
print(Tensor([True, True, False, False]).bitwise_or(Tensor([True, False, True, False])).numpy())
```
"""
self._check_dtype()
return self._binop(Ops.OR, x, reverse)
def bitwise_xor(self, x, reverse=False):
"""
Computes bitwise xor of `self` and `x`.
Equivalent to `self ^ x`.
Supports broadcasting to a common shape, type promotion, and integer, boolean inputs.
```python exec="true" source="above" session="tensor" result="python"
print(Tensor([-1, -2, 3]).bitwise_xor(Tensor([1, 0, 3])).numpy())
```
```python exec="true" source="above" session="tensor" result="python"
print(Tensor([True, True, False, False]).bitwise_xor(Tensor([True, False, True, False])).numpy())
```
"""
self._check_dtype()
return self._binop(Ops.XOR, x, reverse)
def idiv(self, x, reverse=False):
"""
Divides `self` by `x`.