dev, doc, test(hansbug): complete isfinite, isinf, isnan

test/torch/test_funcs.py

@@ -644,3 +644,45 @@ class TestTorchFuncs:
             'a': [[19, 10], [43, 26]],
             'b': {'x': [[44, 32], [80, 59]]},
         })).all()
+
+    def test_isfinite(self):
+        t1 = ttorch.isfinite(torch.tensor([1, float('inf'), 2, float('-inf'), float('nan')]))
+        assert isinstance(t1, torch.Tensor)
+        assert (t1 == ttorch.tensor([True, False, True, False, False])).all()
+
+        t2 = ttorch.isfinite(ttorch.tensor({
+            'a': [1, float('inf'), 2, float('-inf'), float('nan')],
+            'b': {'x': [[1, float('inf'), -2], [float('-inf'), 3, float('nan')]]}
+        }))
+        assert (t2 == ttorch.tensor({
+            'a': [True, False, True, False, False],
+            'b': {'x': [[True, False, True], [False, True, False]]},
+        }))
+
+    def test_isinf(self):
+        t1 = ttorch.isinf(torch.tensor([1, float('inf'), 2, float('-inf'), float('nan')]))
+        assert isinstance(t1, torch.Tensor)
+        assert (t1 == ttorch.tensor([False, True, False, True, False])).all()
+
+        t2 = ttorch.isinf(ttorch.tensor({
+            'a': [1, float('inf'), 2, float('-inf'), float('nan')],
+            'b': {'x': [[1, float('inf'), -2], [float('-inf'), 3, float('nan')]]}
+        }))
+        assert (t2 == ttorch.tensor({
+            'a': [False, True, False, True, False],
+            'b': {'x': [[False, True, False], [True, False, False]]},
+        }))
+
+    def test_isnan(self):
+        t1 = ttorch.isnan(torch.tensor([1, float('inf'), 2, float('-inf'), float('nan')]))
+        assert isinstance(t1, torch.Tensor)
+        assert (t1 == ttorch.tensor([False, False, False, False, True])).all()
+
+        t2 = ttorch.isnan(ttorch.tensor({
+            'a': [1, float('inf'), 2, float('-inf'), float('nan')],
+            'b': {'x': [[1, float('inf'), -2], [float('-inf'), 3, float('nan')]]}
+        }))
+        assert (t2 == ttorch.tensor({
+            'a': [False, False, False, False, True],
+            'b': {'x': [[False, False, False], [False, False, True]]},
+        })).all()

test/torch/test_tensor.py

@@ -281,3 +281,45 @@ class TestTorchTensor:
             'a': [[19, 10], [43, 26]],
             'b': {'x': [[44, 32], [80, 59]]},
         })).all()
+
+    def test_isfinite(self):
+        t1 = torch.tensor([1, float('inf'), 2, float('-inf'), float('nan')]).isfinite()
+        assert isinstance(t1, torch.Tensor)
+        assert (t1 == ttorch.tensor([True, False, True, False, False])).all()
+
+        t2 = ttorch.tensor({
+            'a': [1, float('inf'), 2, float('-inf'), float('nan')],
+            'b': {'x': [[1, float('inf'), -2], [float('-inf'), 3, float('nan')]]}
+        }).isfinite()
+        assert (t2 == ttorch.tensor({
+            'a': [True, False, True, False, False],
+            'b': {'x': [[True, False, True], [False, True, False]]},
+        }))
+
+    def test_isinf(self):
+        t1 = torch.tensor([1, float('inf'), 2, float('-inf'), float('nan')]).isinf()
+        assert isinstance(t1, torch.Tensor)
+        assert (t1 == ttorch.tensor([False, True, False, True, False])).all()
+
+        t2 = ttorch.tensor({
+            'a': [1, float('inf'), 2, float('-inf'), float('nan')],
+            'b': {'x': [[1, float('inf'), -2], [float('-inf'), 3, float('nan')]]}
+        }).isinf()
+        assert (t2 == ttorch.tensor({
+            'a': [False, True, False, True, False],
+            'b': {'x': [[False, True, False], [True, False, False]]},
+        }))
+
+    def test_isnan(self):
+        t1 = torch.tensor([1, float('inf'), 2, float('-inf'), float('nan')]).isnan()
+        assert isinstance(t1, torch.Tensor)
+        assert (t1 == ttorch.tensor([False, False, False, False, True])).all()
+
+        t2 = ttorch.tensor({
+            'a': [1, float('inf'), 2, float('-inf'), float('nan')],
+            'b': {'x': [[1, float('inf'), -2], [float('-inf'), 3, float('nan')]]}
+        }).isnan()
+        assert (t2 == ttorch.tensor({
+            'a': [False, False, False, False, True],
+            'b': {'x': [[False, False, False], [False, False, True]]},
+        })).all()

treetensor/torch/funcs.py

@@ -22,6 +22,7 @@ __all__ = [
     'eq', 'ne', 'lt', 'le', 'gt', 'ge',
     'equal', 'tensor', 'clone',
     'dot', 'matmul', 'mm',
+    'isfinite', 'isinf', 'isnan',
 ]
 
 func_treelize = post_process(post_process(args_mapping(
@@ -954,3 +955,87 @@ def mm(input, mat2, *args, **kwargs):
                               [80, 59]])
     """
     return torch.mm(input, mat2, *args, **kwargs)
+
+
+# noinspection PyShadowingBuiltins
+@doc_from(torch.isfinite)
+@func_treelize()
+def isfinite(input):
+    """
+    In ``treetensor``, you can get a tree of new tensors with boolean elements
+    representing if each element is `finite` or not.
+
+    Examples::
+
+        >>> import torch
+        >>> import treetensor.torch as ttorch
+        >>> ttorch.isfinite(torch.tensor([1, float('inf'), 2, float('-inf'), float('nan')]))
+        tensor([ True, False,  True, False, False])
+
+        >>> ttorch.isfinite(ttorch.tensor({
+        ...     'a': [1, float('inf'), 2, float('-inf'), float('nan')],
+        ...     'b': {'x': [[1, float('inf'), -2], [float('-inf'), 3, float('nan')]]}
+        ... }))
+        <Tensor 0x7fb782a15970>
+        ├── a --> tensor([ True, False,  True, False, False])
+        └── b --> <Tensor 0x7fb782a1e040>
+            └── x --> tensor([[ True, False,  True],
+                              [False,  True, False]])
+    """
+    return torch.isfinite(input)
+
+
+# noinspection PyShadowingBuiltins
+@doc_from(torch.isinf)
+@func_treelize()
+def isinf(input):
+    """
+    In ``treetensor``, you can test if each element of ``input``
+    is infinite (positive or negative infinity) or not.
+
+    Examples::
+
+        >>> import torch
+        >>> import treetensor.torch as ttorch
+        >>> ttorch.isinf(torch.tensor([1, float('inf'), 2, float('-inf'), float('nan')]))
+        tensor([False,  True, False,  True, False])
+
+        >>> ttorch.isinf(ttorch.tensor({
+        ...     'a': [1, float('inf'), 2, float('-inf'), float('nan')],
+        ...     'b': {'x': [[1, float('inf'), -2], [float('-inf'), 3, float('nan')]]}
+        ... }))
+        <Tensor 0x7fb782a29b80>
+        ├── a --> tensor([False,  True, False,  True, False])
+        └── b --> <Tensor 0x7fb782a2d1f0>
+            └── x --> tensor([[False,  True, False],
+                              [ True, False, False]])
+    """
+    return torch.isinf(input)
+
+
+# noinspection PyShadowingBuiltins
+@doc_from(torch.isnan)
+@func_treelize()
+def isnan(input):
+    """
+    In ``treetensor``, you get a tree of new tensors with boolean elements representing
+    if each element of ``input`` is NaN or not
+
+    Examples::
+
+        >>> import torch
+        >>> import treetensor.torch as ttorch
+        >>> ttorch.isnan(torch.tensor([1, float('inf'), 2, float('-inf'), float('nan')]))
+        tensor([False, False, False, False,  True])
+
+        >>> ttorch.isnan(ttorch.tensor({
+        ...     'a': [1, float('inf'), 2, float('-inf'), float('nan')],
+        ...     'b': {'x': [[1, float('inf'), -2], [float('-inf'), 3, float('nan')]]}
+        ... }))
+        <Tensor 0x7fb782a2d0a0>
+        ├── a --> tensor([False, False, False, False,  True])
+        └── b --> <Tensor 0x7fb782a29d90>
+            └── x --> tensor([[False, False, False],
+                              [False, False,  True]])
+    """
+    return torch.isnan(input)

treetensor/torch/tensor.py

@@ -293,3 +293,27 @@ class Tensor(Torch, metaclass=clsmeta(_to_tensor, allow_dict=True)):
         See :func:`treetensor.torch.matmul`.
         """
         return self.matmul(tensor2, *args, **kwargs)
+
+    @doc_from(torch.Tensor.isfinite)
+    @method_treelize()
+    def isfinite(self):
+        """
+        See :func:`treetensor.torch.isfinite`.
+        """
+        return self.isfinite()
+
+    @doc_from(torch.Tensor.isinf)
+    @method_treelize()
+    def isinf(self):
+        """
+        See :func:`treetensor.torch.isinf`.
+        """
+        return self.isinf()
+
+    @doc_from(torch.Tensor.isnan)
+    @method_treelize()
+    def isnan(self):
+        """
+        See :func:`treetensor.torch.isnan`.
+        """
+        return self.isnan()