Implement torch.isnan (#5273)

tunz · soumith · commit cf71385ec91a · 2018-02-19T19:46:35.000-05:00
* Implement torch.isnan

* Simple python implementation

* Fix typo
diff --git a/docs/source/torch.rst b/docs/source/torch.rst
@@ -152,6 +152,7 @@ Comparison Ops
 .. autofunction:: equal
 .. autofunction:: ge
 .. autofunction:: gt
+.. autofunction:: isnan
 .. autofunction:: kthvalue
 .. autofunction:: le
 .. autofunction:: lt
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -3032,6 +3032,10 @@ def test_logical(self):
         self.assertEqual(neqs.sum(), xne.sum(), 0)
         self.assertEqual(x.nelement(), all.sum())
 
+    def test_isnan(self):
+        x = torch.Tensor([1, float('nan'), 2])
+        self.assertEqual(torch.isnan(x), torch.ByteTensor([0, 1, 0]))
+
     def test_RNGState(self):
         state = torch.get_rng_state()
         stateCloned = state.clone()
diff --git a/torch/functional.py b/torch/functional.py
@@ -5,7 +5,7 @@
 
 __all__ = [
     'split', 'chunk', 'empty_like', 'stack', 'unbind', 'btriunpack', 'matmul', 'det', 'stft',
-    'hann_window', 'hamming_window', 'bartlett_window', 'where',
+    'hann_window', 'hamming_window', 'bartlett_window', 'where', 'isnan'
 ]
 
 
@@ -514,3 +514,25 @@ def where(condition, x, y):
     # the parameter order is changed here; the functional order is the same as numpy; the
     # method follows the usual torch mask semantics of x.fn(mask, y)
     return torch._C._VariableBase.where(x, condition, y)
+
+
+def isnan(tensor):
+    r"""Returns a new tensor with boolean elements representing if each element is NaN or not.
+
+    Arguments:
+        tensor (Tensor): A tensor to check
+
+    Returns:
+        Tensor: A ``torch.ByteTensor`` containing a 1 at each location of NaN elements.
+
+    Example::
+
+        >>> torch.isnan(torch.Tensor([1, float('nan'), 2]))
+         0
+         1
+         0
+        [torch.ByteTensor of size 3]
+    """
+    if not torch.is_tensor(tensor):
+        raise ValueError("The argument is not a tensor")
+    return tensor != tensor
