nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

triu tril _tri to mixin (#15857)

Browse Source
This commit is contained in:
chenyu
2026-04-21 17:10:55 -04:00
committed by GitHub
parent e36ff22538
commit 75ee51a446
4 changed files with 66 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/null/test_tensor_uop_mixin.py
+12
View File
@@ -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()
tinygrad/mixin/__init__.py
+50
View File
@@ -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)
tinygrad/mixin/movement.py
+4
View File
@@ -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:
tinygrad/tensor.py
-50
View File
@@ -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.