tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 12:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

cleanup Tensor.triu and Tensor.tril (#2953)

Browse Source 
`.where` does the dtype and shape conversions for 0, no need to use zeros_like
This commit is contained in:
chenyu
2023-12-29 22:27:18 -05:00
committed by GitHub
parent e53b96fdbb
commit 0d6e264c48
1 changed files with 2 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/tensor.py
+2 -4
View File
@@ -666,10 +666,8 @@ class Tensor:
def _tri(r:sint, c:sint, k:int=0, **kwargs) -> Tensor:
assert all_int((r,c)), "does not support symbolic"
return Tensor.arange(r, **kwargs).unsqueeze(1).expand(r,c) <= Tensor.arange(-k, c-k, **kwargs).unsqueeze(0).expand(r,c)
def triu(self, k:int=0) -> Tensor:
return Tensor._tri(self.shape[-2], self.shape[-1], k=k, device=self.device).where(self, Tensor.zeros_like(self))
def tril(self, k:int=0) -> Tensor:
return Tensor._tri(self.shape[-2], self.shape[-1], k=k+1, device=self.device).where(Tensor.zeros_like(self), self)
def triu(self, k:int=0) -> Tensor: return Tensor._tri(self.shape[-2], self.shape[-1], k=k, device=self.device).where(self, 0)
def tril(self, k:int=0) -> Tensor: return Tensor._tri(self.shape[-2], self.shape[-1], k=k+1, device=self.device).where(0, self)
# ***** mlops (unary) *****