nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

move add/mul to MathTrait [pr] (#10741)

Browse Source 
* move add to MathTrait [pr]

* both add and mul
This commit is contained in:
George Hotz
2025-06-09 21:48:55 -07:00
committed by GitHub
parent c28eceaf44
commit 245b1d3a46
2 changed files with 40 additions and 42 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/tensor.py
+3 -40
View File
@@ -186,6 +186,9 @@ class Tensor(MathTrait):
lhs,rhs = self._broadcasted(x, reverse)
return lhs._apply_uop(fxn, rhs)
# _binop is used by MathTrait
def _binop(self, op, x, reverse): return self._apply_broadcasted_uop(lambda *u: UOp.alu(u[0], op, *u[1:]), x, reverse)
def requires_grad_(self, requires_grad=True) -> Tensor:
self.requires_grad = requires_grad
return self
@@ -3490,26 +3493,6 @@ class Tensor(MathTrait):
# broadcast
return x._broadcast_to(out_shape:=_broadcast_shape(x.shape, y.shape)), y._broadcast_to(out_shape)
def add(self, x:Tensor|ConstType, reverse=False) -> Tensor:
"""
Adds `self` and `x`.
Equivalent to `self + x`.
Supports broadcasting to a common shape, type promotion, and integer, float, boolean inputs.
```python exec="true" source="above" session="tensor" result="python"
Tensor.manual_seed(42)
t = Tensor.randn(4)
print(t.numpy())
```
```python exec="true" source="above" session="tensor" result="python"
print(t.add(20).numpy())
```
```python exec="true" source="above" session="tensor" result="python"
print(t.add(Tensor([[2.0], [3.5]])).numpy())
```
"""
return self._apply_broadcasted_uop(UOp.add, x, reverse)
def sub(self, x:Tensor|ConstType, reverse=False) -> Tensor:
"""
Subtracts `x` from `self`.
@@ -3531,26 +3514,6 @@ class Tensor(MathTrait):
a, b = self._broadcasted(x, reverse)
return a + (-b)
def mul(self, x:Tensor|ConstType, reverse=False) -> Tensor:
"""
Multiplies `self` and `x`.
Equivalent to `self * x`.
Supports broadcasting to a common shape, type promotion, and integer, float, boolean inputs.
```python exec="true" source="above" session="tensor" result="python"
Tensor.manual_seed(42)
t = Tensor.randn(4)
print(t.numpy())
```
```python exec="true" source="above" session="tensor" result="python"
print(t.mul(3).numpy())
```
```python exec="true" source="above" session="tensor" result="python"
print(t.mul(Tensor([[-1.0], [2.0]])).numpy())
```
"""
return self._apply_broadcasted_uop(UOp.mul, x, reverse)
def idiv(self, x:Tensor|ConstType, reverse=False) -> Tensor:
"""
Divides `self` by `x`.
tinygrad/uop/mathtraits.py
+37 -2
View File
@@ -14,8 +14,43 @@ 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 add(self, x, reverse=False): return self._binop(Ops.ADD, x, reverse)
def mul(self, x, reverse=False): return self._binop(Ops.MUL, x, reverse)
def add(self, x, reverse=False):
"""
Adds `self` and `x`.
Equivalent to `self + x`.
Supports broadcasting to a common shape, type promotion, and integer, float, boolean inputs.
```python exec="true" source="above" session="tensor" result="python"
Tensor.manual_seed(42)
t = Tensor.randn(4)
print(t.numpy())
```
```python exec="true" source="above" session="tensor" result="python"
print(t.add(20).numpy())
```
```python exec="true" source="above" session="tensor" result="python"
print(t.add(Tensor([[2.0], [3.5]])).numpy())
```
"""
return self._binop(Ops.ADD, x, reverse)
def mul(self, x, reverse=False):
"""
Multiplies `self` and `x`.
Equivalent to `self * x`.
Supports broadcasting to a common shape, type promotion, and integer, float, boolean inputs.
```python exec="true" source="above" session="tensor" result="python"
Tensor.manual_seed(42)
t = Tensor.randn(4)
print(t.numpy())
```
```python exec="true" source="above" session="tensor" result="python"
print(t.mul(3).numpy())
```
```python exec="true" source="above" session="tensor" result="python"
print(t.mul(Tensor([[-1.0], [2.0]])).numpy())
```
"""
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)