diff --git a/test/null/test_tensor_uop_mixin.py b/test/null/test_tensor_uop_mixin.py index e8ef40396d..286dea8ae6 100644 --- a/test/null/test_tensor_uop_mixin.py +++ b/test/null/test_tensor_uop_mixin.py @@ -167,6 +167,18 @@ class TestTensorUOpCreation(unittest.TestCase): self.assertIs(_strip_unique(Tensor.eye(3).uop), _strip_unique(UOp.eye(3))) def test_eye_rect(self): self.assertIs(_strip_unique(Tensor.eye(2, 4).uop), _strip_unique(UOp.eye(2, 4))) + def test_triu(self): + t = _t(3, 4) + self.assertIs(_strip_unique(t.triu().uop), _strip_unique(t.uop.triu())) + def test_triu_diagonal(self): + t = _t(3, 4) + self.assertIs(_strip_unique(t.triu(diagonal=1).uop), _strip_unique(t.uop.triu(diagonal=1))) + def test_tril(self): + t = _t(3, 4) + self.assertIs(_strip_unique(t.tril().uop), _strip_unique(t.uop.tril())) + def test_tril_diagonal(self): + t = _t(3, 4) + self.assertIs(_strip_unique(t.tril(diagonal=-1).uop), _strip_unique(t.uop.tril(diagonal=-1))) if __name__ == "__main__": unittest.main() diff --git a/tinygrad/mixin/__init__.py b/tinygrad/mixin/__init__.py index 95b52ded89..0108a7b575 100644 --- a/tinygrad/mixin/__init__.py +++ b/tinygrad/mixin/__init__.py @@ -141,6 +141,56 @@ class OpMixin(ElementwiseMixin, ReduceMixin): out_dtype = to_dtype(dtype) if dtype is not None else dtypes.default_float return cls.arange(n, device=device).unsqueeze(-1).eq(cls.arange(m_, device=device)).cast(out_dtype) + @classmethod + def _tri(cls, r:sint, c:sint, diagonal=0, device:str|tuple[str, ...]|None=None) -> Self: + return cls.arange(r, device=device).unsqueeze(-1) + diagonal <= cls.arange(c, device=device) + + def triu(self, diagonal:sint=0) -> Self: + """ + Returns the upper triangular part of the tensor, the other elements are set to 0. + + The argument `diagonal` determines which diagonal is on the boundary. `diagonal = 0` means the main diagonal. + Positive `diagonal` means above the main diagonal, and negative `diagonal` means below the main diagonal. + + ```python exec="true" source="above" session="tensor" result="python" + t = Tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) + print(t.numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + print(t.triu(diagonal=0).numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + print(t.triu(diagonal=1).numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + print(t.triu(diagonal=-1).numpy()) + ``` + """ + return self._tri(self.shape[-2], self.shape[-1], diagonal, self.device).where(self, self.zeros_like()) + + def tril(self, diagonal:sint=0) -> Self: + """ + Returns the lower triangular part of the tensor, the other elements are set to 0. + + The argument `diagonal` determines which diagonal is on the boundary. `diagonal = 0` means the main diagonal. + Positive `diagonal` means above the main diagonal, and negative `diagonal` means below the main diagonal. + + ```python exec="true" source="above" session="tensor" result="python" + t = Tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) + print(t.numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + print(t.tril(diagonal=0).numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + print(t.tril(diagonal=1).numpy()) + ``` + ```python exec="true" source="above" session="tensor" result="python" + print(t.tril(diagonal=-1).numpy()) + ``` + """ + return self._tri(self.shape[-2], self.shape[-1], diagonal+1, self.device).where(self.zeros_like(), self) + def _pad_constant(self, pX, value:float) -> Self: # shrink first for negative pads, then pad with only non-negative values pX = tuple((0, 0) if p is None else p for p in pX) diff --git a/tinygrad/mixin/movement.py b/tinygrad/mixin/movement.py index 6c117c8556..2f4853ee07 100644 --- a/tinygrad/mixin/movement.py +++ b/tinygrad/mixin/movement.py @@ -18,6 +18,10 @@ class MovementMixin: def shape(self) -> tuple[sint, ...]: raise NotImplementedError + @property + def device(self) -> str|tuple[str, ...]: + raise NotImplementedError + # great functions you get! @property def ndim(self) -> int: diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 4b4b8568bf..28cafda951 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -1643,56 +1643,6 @@ class Tensor(OpMixin): values, indices = self._inverse().cummax(axis) return values._inverse(), indices - @staticmethod - def _tri(r:sint, c:sint, diagonal=0, device=None, requires_grad:bool|None=None) -> Tensor: - return (Tensor.arange(r, device=device).unsqueeze(-1) + diagonal <= Tensor.arange(c, device=device)).requires_grad_(requires_grad) - - def triu(self, diagonal:sint=0) -> Tensor: - """ - Returns the upper triangular part of the tensor, the other elements are set to 0. - - The argument `diagonal` determines which diagonal is on the boundary. `diagonal = 0` means the main diagonal. - Positive `diagonal` means above the main diagonal, and negative `diagonal` means below the main diagonal. - - ```python exec="true" source="above" session="tensor" result="python" - t = Tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) - print(t.numpy()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - print(t.triu(diagonal=0).numpy()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - print(t.triu(diagonal=1).numpy()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - print(t.triu(diagonal=-1).numpy()) - ``` - """ - return Tensor._tri(self.shape[-2], self.shape[-1], diagonal=diagonal, device=self.device).where(self, self.zeros_like()) - - def tril(self, diagonal:sint=0) -> Tensor: - """ - Returns the lower triangular part of the tensor, the other elements are set to 0. - - The argument `diagonal` determines which diagonal is on the boundary. `diagonal = 0` means the main diagonal. - Positive `diagonal` means above the main diagonal, and negative `diagonal` means below the main diagonal. - - ```python exec="true" source="above" session="tensor" result="python" - t = Tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]]) - print(t.numpy()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - print(t.tril(diagonal=0).numpy()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - print(t.tril(diagonal=1).numpy()) - ``` - ```python exec="true" source="above" session="tensor" result="python" - print(t.tril(diagonal=-1).numpy()) - ``` - """ - return Tensor._tri(self.shape[-2], self.shape[-1], diagonal=diagonal+1, device=self.device).where(self.zeros_like(), self) - def interpolate(self, size:tuple[int, ...], mode:str="linear", align_corners:bool=False) -> Tensor: """ Downsamples or Upsamples to the input `size`, accepts 0 to N batch dimensions.