Support scalars by kartik4949
MODIFY
test/test_ops.py
test/test_ops.py
@@ -66,9 +66,14 @@ class TestOps(unittest.TestCase):
def test_logsoftmax(self):
helper_test_op([(45,65)], lambda x: torch.nn.LogSoftmax(dim=1)(x), Tensor.logsoftmax, atol=1e-7, grad_atol=1e-7, gpu=self.gpu)
def test_tanh(self):
- helper_test_op([(45,65)], lambda x: x.tanh(), Tensor.tanh, atol=1e-6, grad_atol=1e-6)
+ helper_test_op([(45,65)], lambda x: x.tanh(), Tensor.tanh, atol=1e-6, grad_atol=1e-6, gpu=self.gpu)
def test_topo_sort(self):
- helper_test_op([(45,65)], lambda x: (x+x)*x, lambda x: x.add(x).mul(x), atol=1e-6, grad_atol=1e-6)
+ helper_test_op([(45,65)], lambda x: (x+x)*x, lambda x: x.add(x).mul(x), atol=1e-6, grad_atol=1e-6, gpu=self.gpu)
+
+ def test_scalar_mul(self):
+ helper_test_op([(45,65)], lambda x: x*2, lambda x: x*2, gpu=self.gpu)
+ def test_scalar_rmul(self):
+ helper_test_op([(45,65)], lambda x: 2*x, lambda x: 2*x, gpu=self.gpu)
def test_broadcast_full(self):
for torch_op, tinygrad_op in [(torch.add, Tensor.add), (torch.sub, Tensor.sub), (torch.mul, Tensor.mul),

MODIFY
tinygrad/ops_cpu.py
tinygrad/ops_cpu.py
@@ -6,8 +6,8 @@ from .tensor import Function, register
# ************* basic ops *************
def unbroadcast(out, in_sh):
# adjoint operation to broadcast is sum. Need to sum all axis with 1 = in_sh[i] < out.shape[i]
- sum_axis = [i for i in range(len(in_sh)) if in_sh[i]==1 and out.shape[i]>1]
- return out.sum(axis=tuple(sum_axis)).reshape(in_sh)
+ sum_axis = tuple([i for i in range(len(in_sh)) if in_sh[i]==1 and out.shape[i]>1]) if in_sh != (1,) else None
+ return out.sum(axis=sum_axis).reshape(in_sh)
class Add(Function):
@staticmethod

MODIFY
tinygrad/ops_gpu.py
tinygrad/ops_gpu.py
@@ -170,7 +170,7 @@ def reduce_op(ctx, code, code2, inp, axis=None):
return ret
def unbroadcast(ctx, out, in_sh):
- sum_axis = [i for i in range(len(in_sh)) if in_sh[i]==1 and out.shape[i]>1]
+ sum_axis = [i for i in range(len(in_sh)) if in_sh[i]==1 and out.shape[i]>1] if in_sh != (1,) else None
return reduce_op(ctx, "out += a", "out", out, sum_axis)
# ***** now for the ops themselves *****

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -197,24 +197,19 @@ class Tensor:
def mean(self, axis=None):
out = self.sum(axis=axis)
coeff = np.prod(out.shape)/np.prod(self.shape)
- div = Tensor(coeff+np.zeros(out.shape, dtype=self.dtype), gpu=self.gpu, requires_grad=False)
- return out.mul(div)
+ return out * coeff
def sqrt(self):
- root = Tensor(np.zeros(self.shape, dtype=self.dtype)+0.5, gpu=self.gpu, requires_grad=False)
- return self.pow(root)
+ return self.pow(0.5)
def div(self, y):
- root = Tensor(np.zeros(self.shape, dtype=self.dtype)-1, gpu=self.gpu, requires_grad=False)
- return self.mul(y.pow(root))
+ return self * (y ** -1.0)
def swish(self):
- return self.mul(self.sigmoid())
+ return self * self.sigmoid()
def tanh(self):
- t2 = Tensor(np.zeros(self.shape, dtype=self.dtype)+2, gpu=self.gpu, requires_grad=False)
- t1 = Tensor(np.zeros(self.shape, dtype=self.dtype)+1, gpu=self.gpu, requires_grad=False)
- return self.mul(t2).sigmoid().mul(t2) - t1 # 2*sigmoid(2*x)-1
+ return 2.0 * ((2.0 * self).sigmoid()) - 1.0
# An instantiation of the Function is the Context
class Function:
@@ -249,13 +244,17 @@ def register(name, fxn, gpu=False):
else:
Tensor.ops_cpu[name] = fxn
def dispatch(*x, **kwargs):
+ assert isinstance(x[0], Tensor)
+ x = [Tensor(np.array([arg], dtype=x[0].dtype), gpu=x[0].gpu, requires_grad=False) if not isinstance(arg, Tensor) else arg for arg in x]
f = (Tensor.ops_gpu if x[0].gpu else Tensor.ops_cpu)[name]
f.cl_ctx, f.cl_queue = cl_ctx, cl_queue
return f.apply(f, *x, **kwargs)
setattr(Tensor, name, dispatch)
- if name in ['add', 'sub', 'mul', 'div', 'pow']:
+ # TODO: div is a second class op, so it doesn't work here
+ if name in ['add', 'sub', 'mul', 'pow']:
setattr(Tensor, f"__{name}__", dispatch)
setattr(Tensor, f"__i{name}__", lambda self,x: self.assign(dispatch(self,x)))
+ setattr(Tensor, f"__r{name}__", lambda self,x: dispatch(self,x))
# this registers all the operations
import tinygrad.ops_cpu

Support scalars by kartik4949