diff --git a/docs/tensor/ops.md b/docs/tensor/ops.md index faaea2ff3c..5fbf376ff8 100644 --- a/docs/tensor/ops.md +++ b/docs/tensor/ops.md @@ -1,6 +1,7 @@ ## Reduce ::: tinygrad.Tensor.sum +::: tinygrad.Tensor.prod ::: tinygrad.Tensor.max ::: tinygrad.Tensor.min ::: tinygrad.Tensor.any diff --git a/extra/onnx_ops.py b/extra/onnx_ops.py index 82db4d6f08..ee08726269 100644 --- a/extra/onnx_ops.py +++ b/extra/onnx_ops.py @@ -56,7 +56,6 @@ Softmax = {1: Softmax_1, 13: Softmax_13} # Softmax default axis changed def LogSoftmax(x: Tensor, axis=-1): return x.log_softmax(axis) def Clip(x: Tensor, min=None, max=None): return x.clip(float('-inf') if min is None else min, float('inf') if max is None else max).cast(x.dtype) -# NOTE ReduceProd would require a new llop def _axes(axes, noop_with_empty_axes): if axes is not None and not (isinstance(axes, Tensor) and axes.shape == (0,)): return to_python_const(axes) return [] if noop_with_empty_axes else None @@ -65,6 +64,7 @@ def ReduceMin(data: Tensor, axes=None, keepdims=1, noop_with_empty_axes=0): retu def ReduceSum(data: Tensor, axes=None, keepdims=1, noop_with_empty_axes=0): return data.sum(_axes(axes, noop_with_empty_axes), keepdim=keepdims) def ReduceMean(data: Tensor, axes=None, keepdims=1, noop_with_empty_axes=0): return data.mean(_axes(axes, noop_with_empty_axes), keepdim=keepdims) def ReduceSumSquare(data: Tensor, axes=None, keepdims=1, noop_with_empty_axes=0): return ReduceSum(data.square(), axes, keepdims, noop_with_empty_axes) +def ReduceProd(data: Tensor, axes=None, keepdims=1, noop_with_empty_axes=0): return data.prod(_axes(axes, noop_with_empty_axes), keepdim=keepdims) def ReduceL1(data: Tensor, axes=None, keepdims=1, noop_with_empty_axes=0): return ReduceSum(data.abs(), axes, keepdims, noop_with_empty_axes) def ReduceL2(data: Tensor, axes=None, keepdims=1, noop_with_empty_axes=0): return ReduceSumSquare(data, axes, keepdims, noop_with_empty_axes).sqrt() def ReduceLogSum(data: Tensor, axes=None, keepdims=1, noop_with_empty_axes=0): return ReduceSum(data, axes, keepdims, noop_with_empty_axes).log() diff --git a/test/external/external_test_onnx_backend.py b/test/external/external_test_onnx_backend.py index 10026c8ea5..232cecc5e1 100644 --- a/test/external/external_test_onnx_backend.py +++ b/test/external/external_test_onnx_backend.py @@ -40,9 +40,6 @@ class TinygradBackend(Backend): backend_test = onnx.backend.test.BackendTest(TinygradBackend, __name__) -# no support for reduce with multiply (needs llop) -backend_test.exclude('test_reduce_prod_*') - # TODO figure out why it's returning wrong values, geohotstan's uneducated guess is it's due to imprecision from float64 (double) -> float32 # see Type Constraints: https://onnx.ai/onnx/operators/onnx_aionnxpreviewtraining_Adam.html#type-constraints backend_test.exclude('test_adam_multiple_cpu') diff --git a/test/test_const_folding.py b/test/test_const_folding.py index f9016a1e32..2fa551b136 100644 --- a/test/test_const_folding.py +++ b/test/test_const_folding.py @@ -150,6 +150,14 @@ class TestReduceOpsConstFolding(unittest.TestCase): _check_ast_count(1, Tensor.ones(4).pad(((1, 1),)).exp().sum()) np.testing.assert_allclose(Tensor.ones(4).pad(((1, 1),)).exp().sum().numpy(), 4 * math.e + 2) + def test_const_prod(self): + _check_ast_count(0, Tensor.full((2, 3), fill_value=2).prod()) + np.testing.assert_equal(Tensor.full((2, 3), fill_value=2).prod().numpy(), 2**(2*3)) + _check_ast_count(0, Tensor.full((4, 5, 6), fill_value=2).prod(axis=0)) + np.testing.assert_equal(Tensor.full((4, 5, 6), fill_value=2).prod(axis=0).numpy(), np.full((5, 6), 2**4)) + _check_ast_count(0, Tensor(4).prod()) + np.testing.assert_equal(Tensor(4).prod().numpy(), 4) + def test_const_max(self): _check_ast_count(0, Tensor.ones(4, 5, 6).max()) np.testing.assert_equal(Tensor.ones(4, 5, 6).max().numpy(), 1) diff --git a/test/test_dtype.py b/test/test_dtype.py index cb26729e1c..42577acdbf 100644 --- a/test/test_dtype.py +++ b/test/test_dtype.py @@ -621,6 +621,13 @@ class TestAutoCastType(unittest.TestCase): assert t.sum(acc_dtype=dtypes.float32).dtype == dtypes.float32 np.testing.assert_allclose(t.sum(acc_dtype=dtypes.float32).numpy(), 80000) + def test_prod_acc_dtype(self): + t = Tensor([100, 200], dtype=dtypes.int32) + assert t.prod().dtype == dtypes.int32 + np.testing.assert_allclose(t.prod().numpy(), 20000) + assert t.prod(acc_dtype=dtypes.float32).dtype == dtypes.float32 + np.testing.assert_allclose(t.prod(acc_dtype=dtypes.float32).numpy(), 20000) + def test_mean(self): assert (Tensor([0, 1], dtype=dtypes.bool)).mean().dtype == dtypes.float32 assert (Tensor([0, 1], dtype=dtypes.int8)).mean().dtype == dtypes.float32 diff --git a/test/test_ops.py b/test/test_ops.py index 9cc07ddf55..8e5de5fd9a 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -878,6 +878,7 @@ class TestOps(unittest.TestCase): helper_test_op([(3,4,5,6)], lambda x: x.sum(axis=(0,2))) helper_test_op([(3,4,5,6)], lambda x: x.sum(axis=(1,2))) helper_test_op([(3,4,5,6)], lambda x: x.sum(axis=1)) + helper_test_op([(3,4,5,6)], lambda x: x.sum(axis=1, keepdim=True)) helper_test_op([()], lambda x: x.sum()) helper_test_op([()], lambda x: x.sum(0)) helper_test_op([()], lambda x: x.sum(-1)) @@ -891,6 +892,15 @@ class TestOps(unittest.TestCase): helper_test_op([(4, 0)], lambda x: x.sum(axis=(1,))) helper_test_op([(4, 0)], lambda x: x.sum(axis=(0,1))) + def test_prod(self): + helper_test_op(None, lambda x: x.prod(), vals=[[1.0, 2.0, 3.0]]) + helper_test_op([(3,4,5,6)], lambda x: x.prod(dim=3), lambda x: x.prod(axis=3)) + helper_test_op([(3,4,5,6)], lambda x: x.prod(dim=1), lambda x: x.prod(axis=1)) + helper_test_op([(3,4,5,6)], lambda x: x.prod(dim=1, keepdim=True), lambda x: x.prod(axis=1, keepdim=True)) + helper_test_op([()], lambda x: x.prod()) + helper_test_op([()], lambda x: x.prod(0)) + helper_test_op([()], lambda x: x.prod(-1)) + def test_min(self): helper_test_op([(3,3)], lambda x: x.min()) helper_test_op([(45,3)], lambda x: x.min()) diff --git a/tinygrad/codegen/lowerer.py b/tinygrad/codegen/lowerer.py index ec32326ff4..204a5a16ab 100644 --- a/tinygrad/codegen/lowerer.py +++ b/tinygrad/codegen/lowerer.py @@ -121,7 +121,7 @@ class IndependentLowerer: return UOp(UOps.EXPAND, x.dtype, tuple(UOp(UOps.GEP, x.dtype, (ret,), i) for i in range(wmma_sz[2])), arg=upcast_axes[2]) # NOTE: always using ridxs is fine here reduce_range, reduce_expand = partition([self.ridxs[i] for i in x.arg[1]], lambda y: y.op is UOps.RANGE) - alu_op = {ReduceOps.SUM:BinaryOps.ADD, ReduceOps.MAX:BinaryOps.MAX}[cast(ReduceOps, x.arg[0])] + alu_op = {ReduceOps.SUM:BinaryOps.ADD, ReduceOps.PROD:BinaryOps.MUL, ReduceOps.MAX:BinaryOps.MAX}[cast(ReduceOps, x.arg[0])] ret = in_uops[0] if len(contract_axis:=flatten(x.arg for x in reduce_expand)): ret = UOp(UOps.CONTRACT, cast(DType, x.dtype).vec(prod(x[1] for x in contract_axis)), (ret,), tuple(contract_axis)) diff --git a/tinygrad/function.py b/tinygrad/function.py index 925717939e..27a3113586 100644 --- a/tinygrad/function.py +++ b/tinygrad/function.py @@ -151,6 +151,14 @@ class Sum(Function): def backward(self, grad_output:LazyBuffer) -> LazyBuffer: return grad_output.expand(self.input_shape) +class Prod(Function): + def forward(self, x:LazyBuffer, axis:Tuple[int, ...]) -> LazyBuffer: + self.x, self.ret = x, x.r(ReduceOps.PROD, axis) + return self.ret + + def backward(self, grad_output:LazyBuffer) -> LazyBuffer: + return grad_output.e(BinaryOps.MUL, self.ret).expand(self.x.shape).e(BinaryOps.MUL, self.x.e(UnaryOps.RECIP)) + class Max(Function): def forward(self, x:LazyBuffer, axis:Tuple[int, ...]) -> LazyBuffer: self.x, self.ret, self.axis = x, x.r(ReduceOps.MAX, axis), axis diff --git a/tinygrad/lazy.py b/tinygrad/lazy.py index d8d44dfaef..0c99745115 100644 --- a/tinygrad/lazy.py +++ b/tinygrad/lazy.py @@ -173,12 +173,15 @@ class LazyBuffer: def r(self, op:ReduceOps, axis:Tuple[int, ...]) -> LazyBuffer: new_shape = self.st.reduce(axis) # TODO: this logic should move to the scheduler - if 0 in self.shape and 0 not in new_shape: return self.const({ReduceOps.SUM: 0.0, ReduceOps.MAX: dtypes.min(self.dtype)}[op], new_shape) + if 0 in self.shape and 0 not in new_shape: + return self.const({ReduceOps.SUM: 0.0, ReduceOps.PROD: 1.0, ReduceOps.MAX: dtypes.min(self.dtype)}[op], new_shape) # const folding # TODO: fold this for symbolic? if self.is_unrealized_unmasked_const() and all_int(self.shape): - return self.const(self.base.arg * {ReduceOps.SUM: prod(self.shape[i] for i in axis), ReduceOps.MAX: 1}[op], new_shape) + if op is ReduceOps.SUM: return self.const(self.base.arg * prod(self.shape[i] for i in axis), new_shape) + if op is ReduceOps.PROD: return self.const(self.base.arg ** prod(self.shape[i] for i in axis), new_shape) + if op is ReduceOps.MAX: return self.const(self.base.arg, new_shape) # TODO: can we split symbolic shape if the reduce axis is not symbolic? if not SPLIT_REDUCEOP or not all_int(self.shape) or (0 in self.shape) or \ diff --git a/tinygrad/nn/__init__.py b/tinygrad/nn/__init__.py index 512477074e..b7e148169c 100644 --- a/tinygrad/nn/__init__.py +++ b/tinygrad/nn/__init__.py @@ -32,8 +32,8 @@ class BatchNorm: def __init__(self, sz:int, eps=1e-5, affine=True, track_running_stats=True, momentum=0.1): self.eps, self.track_running_stats, self.momentum = eps, track_running_stats, momentum - if affine: self.weight, self.bias = Tensor.ones(sz), Tensor.zeros(sz) - else: self.weight, self.bias = None, None + self.weight: Optional[Tensor] = Tensor.ones(sz) if affine else None + self.bias: Optional[Tensor] = Tensor.zeros(sz) if affine else None self.num_batches_tracked = Tensor.zeros(1, requires_grad=False) if track_running_stats: self.running_mean, self.running_var = Tensor.zeros(sz, requires_grad=False), Tensor.ones(sz, requires_grad=False) diff --git a/tinygrad/ops.py b/tinygrad/ops.py index b7f6d3dd02..7465665720 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -25,7 +25,7 @@ class TernaryOps(Enum): WHERE = auto(); MULACC = auto() # noqa: E702 class ReduceOps(Enum): """A -> B (reduce)""" - SUM = auto(); MAX = auto(); WMMA = auto() # noqa: E702 + SUM = auto(); PROD = auto(); MAX = auto(); WMMA = auto() # noqa: E702 class MetaOps(Enum): EMPTY = auto(); CONST = auto(); COPY = auto(); CONTIGUOUS = auto(); CUSTOM = auto(); ASSIGN = auto(); VIEW = auto() # noqa: E702 Op = Union[UnaryOps, BinaryOps, ReduceOps, MetaOps, TernaryOps] diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 91e83ff082..d5fc7fc58b 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -459,7 +459,7 @@ class Tensor: # ***** creation helper functions ***** @staticmethod - def full(shape:Tuple[sint, ...], fill_value:ConstType, **kwargs): + def full(shape:Tuple[sint, ...], fill_value:ConstType, **kwargs) -> Tensor: """ Creates a tensor with the given shape, filled with the given value. @@ -476,7 +476,7 @@ class Tensor: return Tensor(fill_value, **kwargs).reshape((1, )*len(new_shape := argfix(shape))).expand(new_shape) @staticmethod - def zeros(*shape, **kwargs): + def zeros(*shape, **kwargs) -> Tensor: """ Creates a tensor with the given shape, filled with zeros. @@ -493,7 +493,7 @@ class Tensor: return Tensor.full(argfix(*shape), 0.0, **kwargs) @staticmethod - def ones(*shape, **kwargs): + def ones(*shape, **kwargs) -> Tensor: """ Creates a tensor with the given shape, filled with ones. @@ -510,7 +510,7 @@ class Tensor: return Tensor.full(argfix(*shape), 1.0, **kwargs) @staticmethod - def arange(start, stop=None, step=1, **kwargs): + def arange(start, stop=None, step=1, **kwargs) -> Tensor: """ Returns a 1-D tensor of size `ceil((stop - start) / step)` with values from `[start, stop)`, with spacing between values given by `step`. @@ -542,7 +542,7 @@ class Tensor: return (Tensor.full((math.ceil((stop-start)/step),), step, dtype=dtype, **kwargs)._cumsum() + (start - step)).cast(dtype) @staticmethod - def eye(n:int, m:Optional[int]=None, **kwargs): + def eye(n:int, m:Optional[int]=None, **kwargs) -> Tensor: """ Returns a 2-D tensor with `n` rows and `m` columns, with ones on the diagonal and zeros elsewhere. @@ -560,7 +560,7 @@ class Tensor: if n < 0 or (m is not None and m < 0): raise ValueError(f"cannot have negative {n=}, {m=}") return Tensor.ones((n,1),**kwargs).pad((None,(0,n))).flatten().shrink(((0,n*n),)).reshape(n,n)._slice((None,(0,n if m is None else m))) - def full_like(self, fill_value:ConstType, **kwargs): + def full_like(self, fill_value:ConstType, **kwargs) -> Tensor: """ Creates a tensor with the same shape as `self`, filled with the given value. If `dtype` is not specified, the dtype of `self` is used. @@ -575,7 +575,7 @@ class Tensor: """ return Tensor.full(self.shape, fill_value, dtype=kwargs.pop("dtype", self.dtype), device=kwargs.pop("device", self.device), **kwargs) - def zeros_like(self, **kwargs): + def zeros_like(self, **kwargs) -> Tensor: """ Creates a tensor with the same shape as `self`, filled with zeros. @@ -589,7 +589,7 @@ class Tensor: """ return self.full_like(0, **kwargs) - def ones_like(self, **kwargs): + def ones_like(self, **kwargs) -> Tensor: """ Creates a tensor with the same shape as `self`, filled with ones. @@ -1330,7 +1330,7 @@ class Tensor: def sum(self, axis:Optional[Union[int, Sequence[int]]]=None, keepdim=False, acc_dtype:Optional[DTypeLike]=None): """ - Sums the elements of the tensor along the specified axis or axes. + Returns the sum of the elements of the tensor along the specified axis or axes. You can pass in `axis` and `keepdim` keyword arguments to control the axis along which the maximum is computed and whether the reduced dimensions are retained. @@ -1355,6 +1355,32 @@ class Tensor: ret = self.cast(acc_dtype or sum_acc_dtype(self.dtype))._reduce(F.Sum, axis, keepdim) return ret.cast(self.dtype) if acc_dtype is None and self.dtype in (dtypes.float16, dtypes.bfloat16) else ret + def prod(self, axis:Optional[Union[int, Sequence[int]]]=None, keepdim=False, acc_dtype:Optional[DTypeLike]=None): + """ + Returns the product of the elements of the tensor along the specified axis or axes. + + You can pass in `axis` and `keepdim` keyword arguments to control the axis along + which the maximum is computed and whether the reduced dimensions are retained. + + You can pass in `acc_dtype` keyword argument to control the data type of the accumulation. + If not specified, the accumulation data type is chosen based on the input tensor's data type. + + ```python exec="true" source="above" session="tensor" result="python" + t = Tensor([-1, -2, -3, 1, 2, 3]).reshape(2, 3) + print(t.numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + print(t.prod().numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + print(t.prod(axis=0).numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + print(t.prod(axis=1).numpy()) + ``` + """ + return self.cast(acc_dtype or self.dtype)._reduce(F.Prod, axis, keepdim) + def max(self, axis:Optional[Union[int, Sequence[int]]]=None, keepdim=False): """ Returns the maximum value of the tensor along the specified axis or axes.