diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 328fcc038e..6a94d49951 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -1261,7 +1261,7 @@ class Tensor: # ***** reduce ops ***** - def _reduce(self, fxn:Type[Function], axis:Optional[Union[int, Tuple[int, ...]]]=None, keepdim=False) -> Tensor: + def _reduce(self, fxn:Type[Function], axis:Optional[Union[int, Sequence[int]]]=None, keepdim=False) -> Tensor: if self.ndim == 0: if axis is not None and axis not in [-1, 0]: raise IndexError(f"{axis=} out of range of [-1, 0]") axis = None @@ -1270,7 +1270,7 @@ class Tensor: ret = fxn.apply(self, axis=axis_) return ret if keepdim else ret.reshape(tuple(s for i,s in enumerate(self.shape) if i not in axis_)) - def sum(self, axis=None, keepdim=False, acc_dtype:Optional[DType]=None): + def sum(self, axis:Optional[Union[int, Sequence[int]]]=None, keepdim=False, acc_dtype:Optional[DType]=None): """ Sums the elements of the tensor along the specified axis or axes. @@ -1297,7 +1297,7 @@ 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 max(self, axis=None, keepdim=False): + 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. @@ -1320,7 +1320,7 @@ class Tensor: """ return self._reduce(F.Max, axis, keepdim) - def min(self, axis=None, keepdim=False): + def min(self, axis:Optional[Union[int, Sequence[int]]]=None, keepdim=False): """ Returns the minimum value of the tensor along the specified axis or axes. @@ -1343,7 +1343,7 @@ class Tensor: """ return -((-self).max(axis=axis, keepdim=keepdim)) - def mean(self, axis=None, keepdim=False): + def mean(self, axis:Optional[Union[int, Sequence[int]]]=None, keepdim=False): """ Returns the mean value of the tensor along the specified axis or axes. @@ -1369,7 +1369,7 @@ class Tensor: numerator = self.cast(sum_acc_dtype(self.dtype)).sum(axis=axis, keepdim=keepdim) return numerator.div(prod([si for si, so in zip(self.shape, self.sum(axis=axis, keepdim=True).shape) if si != so])).cast(output_dtype) - def var(self, axis=None, keepdim=False, correction=1): + def var(self, axis:Optional[Union[int, Sequence[int]]]=None, keepdim=False, correction=1): """ Returns the variance of the tensor along the specified axis or axes. @@ -1395,7 +1395,7 @@ class Tensor: n = prod([si for si, so in zip(self.shape, squares.sum(axis=axis, keepdim=True).shape) if si != so]) return squares.sum(axis=axis, keepdim=keepdim).div(max(0, n-correction)) - def std(self, axis=None, keepdim=False, correction=1): + def std(self, axis:Optional[Union[int, Sequence[int]]]=None, keepdim=False, correction=1): """ Returns the standard deviation of the tensor along the specified axis or axes.