Introducing IsInf (pytorch#9169)

bhushan23 · jramseyer · commit c21388f30c01 · 2018-07-30T16:02:15.000-04:00
Summary: torch.isinf - checks element wise +/- inf implements pytorch#9132 Pull Request resolved: pytorch#9169 Reviewed By: SsnL Differential Revision: D8768614 Pulled By: zou3519 fbshipit-source-id: dd1b5f6c976deb421d626e22cdd25500ec04d796
diff --git a/docs/source/torch.rst b/docs/source/torch.rst
@@ -223,6 +223,7 @@ Comparison Ops
 .. autofunction:: equal
 .. autofunction:: ge
 .. autofunction:: gt
+.. autofunction:: isinf
 .. autofunction:: isnan
 .. autofunction:: kthvalue
 .. autofunction:: le
diff --git a/test/test_torch.py b/test/test_torch.py
@@ -4711,6 +4711,10 @@ def test_isnan(self):
         x = torch.Tensor([1, float('nan'), 2])
         self.assertEqual(torch.isnan(x), torch.ByteTensor([0, 1, 0]))
 
+    def test_isinf(self):
+        x = torch.Tensor([1, float('inf'), 2, float('-inf'), float('nan')])
+        self.assertEqual(torch.isinf(x), torch.ByteTensor([0, 1, 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
@@ -8,6 +8,7 @@
     'argmin',
     'btrifact',
     'btriunpack',
+    'isinf',
     'isnan',
     'split',
     'unique',
@@ -135,6 +136,25 @@ def btriunpack(LU_data, LU_pivots, unpack_data=True, unpack_pivots=True):
     return P, L, U
 
 
+def isinf(tensor):
+    r"""Returns a new tensor with boolean elements representing if each element is `+/-INF` or not.
+
+    Arguments:
+        tensor (Tensor): A tensor to check
+
+    Returns:
+        Tensor: A ``torch.ByteTensor`` containing a 1 at each location of `+/-INF` elements and 0 otherwise
+
+    Example::
+
+        >>> torch.isinf(torch.Tensor([1, float('inf'), 2, float('-inf'), float('nan')]))
+        tensor([ 0,  1,  0,  1,  0], dtype=torch.uint8)
+    """
+    if not isinstance(tensor, torch.Tensor):
+        raise ValueError("The argument is not a tensor", str(tensor))
+    return tensor.abs() == float('inf')
+
+
 def isnan(tensor):
     r"""Returns a new tensor with boolean elements representing if each element is `NaN` or not.
 
@@ -150,7 +170,7 @@ def isnan(tensor):
         tensor([ 0,  1,  0], dtype=torch.uint8)
     """
     if not isinstance(tensor, torch.Tensor):
-        raise ValueError("The argument is not a tensor")
+        raise ValueError("The argument is not a tensor", str(tensor))
     return tensor != tensor
 
 
