github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-08 22:48:25 -05:00
Code Issues Packages Projects Releases Wiki Activity

Add isinf and isnan ops to Tensor (#7484)

Browse Source 
* move isinf and isnan to new branch

* sneak a roll documentation fix in

* add to docs

* update test coverage for detect_positive and detect_negative

* add types to isinf args
This commit is contained in:
geohotstan
2024-11-03 00:12:52 +08:00
committed by GitHub
parent 72a9ac27e9
commit 585f3a0f24
3 changed files with 32 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
docs/tensor/elementwise.md
View File

@@ -17,6 +17,8 @@ Elementwise ops operate on a per element basis. They don't change the shape of t
::: tinygrad.Tensor.ceil
::: tinygrad.Tensor.floor
::: tinygrad.Tensor.round
::: tinygrad.Tensor.isinf
::: tinygrad.Tensor.isnan
::: tinygrad.Tensor.lerp
::: tinygrad.Tensor.square
::: tinygrad.Tensor.clamp

8
test/test_ops.py
View File

@@ -339,6 +339,14 @@ class TestOps(unittest.TestCase):
helper_test_op(None, lambda x: x.round(), vals=[[1.499, 1.5, 1.501, 1.0, 2.1, 0.0, -5.0, -2.499, -2.5, -2.501]], forward_only=True)
helper_test_op(None, lambda x: x.round(), vals=[[2.5, -1.5]], forward_only=True)
def test_isinf(self):
val = [float('-inf'), 0., float('inf'), float('nan'), 1.1]
helper_test_op(None, torch.isinf, Tensor.isinf, vals=[val], forward_only=True)
np.testing.assert_equal(Tensor(val).isinf(detect_positive=True, detect_negative=False).numpy(), [False, False, True, False, False])
np.testing.assert_equal(Tensor(val).isinf(detect_positive=False, detect_negative=True).numpy(), [True, False, False, False, False])
def test_isnan(self):
helper_test_op(None, torch.isnan, Tensor.isnan, vals=[[float('-inf'), 0., float('inf'), float('nan'), 1.1]], forward_only=True)
def test_lerp(self):
helper_test_op([(45,35), (45,35), (45,35)], lambda x,y,z: x.lerp(y,z))
helper_test_op(None, lambda x,y,z: x.lerp(y,z), vals=[[1.,2.,3.], [4.,5.,6.], 0.5])

24
tinygrad/tensor.py
View File

@@ -1437,10 +1437,11 @@ class Tensor(SimpleMathTrait): # pylint: disable=abstract-method
The rolling operation is circular, meaning that elements that go beyond the edge are wrapped around to the beginning of the dimension.
```python exec="true" source="above" session="tensor" result="python"
print(Tensor.rand(3, 4, 1).roll(shifts=1, dims=0))
t = Tensor.arange(4)
print(t.roll(shifts=1, dims=0).numpy())
```
```python exec="true" source="above" session="tensor" result="python"
print(Tensor.rand(3, 4, 1).roll(shifts=-1, dims=0))
print(t.roll(shifts=-1, dims=0).numpy())
```
"""
dims, rolled = tuple(self._resolve_dim(d) for d in make_tuple(dims, 1)), self
@@ -2443,6 +2444,25 @@ class Tensor(SimpleMathTrait): # pylint: disable=abstract-method
"""
return ((self > 0) == ((b := self.cast(dtypes.int32) / 2.0).cast(dtypes.int32) == b)).where((self - 0.5).ceil(), (self + 0.5).floor())
def isinf(self:Tensor, detect_positive:bool=True, detect_negative:bool=True):
"""
Checks the tensor element-wise to return True where the element is infinity, otherwise returns False
```python exec="true" source="above" session="tensor" result="python"
print(Tensor([1, float('inf'), 2, float('-inf'), float('nan')]).isinf().numpy())
```
"""
return (self == float("inf")) * detect_positive + (self == float("-inf")) * detect_negative
def isnan(self:Tensor):
"""
Checks the tensor element-wise to return True where the element is NaN, otherwise returns False
```python exec="true" source="above" session="tensor" result="python"
print(Tensor([1, float('inf'), 2, float('-inf'), float('nan')]).isnan().numpy())
```
"""
return self != self
def lerp(self, end: Tensor, weight: Union[Tensor, float]) -> Tensor:
"""
Linearly interpolates between `self` and `end` by `weight`.