diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index a575951f03..92d646e7cf 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -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`. diff --git a/tinygrad/uop/mathtraits.py b/tinygrad/uop/mathtraits.py index 81e4d64111..adc2c9d8fb 100644 --- a/tinygrad/uop/mathtraits.py +++ b/tinygrad/uop/mathtraits.py @@ -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)