Added numpy-like posinf and neginf

Author: Dhruvanshu-Joshi

pytensor/tensor/math.py

@@ -811,6 +811,22 @@ def largest(*args):
         return max(stack(args), axis=0)
 
 
+def isposinf(x):
+    """
+    Return if the input variable has positive infinity element
+
+    """
+    return eq(x, np.inf)
+
+
+def isneginf(x):
+    """
+    Return if the input variable has negative infinity element
+
+    """
+    return eq(x, -np.inf)
+
+
 @scalar_elemwise
 def lt(a, b):
     """a < b"""
@@ -3077,11 +3093,11 @@ def nan_to_num(x, nan=0.0, posinf=None, neginf=None):
     """
     # Replace NaN's with nan keyword
     is_nan = isnan(x)
-    is_pos_inf = eq(x, np.inf)
-    is_neg_inf = eq(x, -np.inf)
+    is_pos_inf = isposinf(x)
+    is_neg_inf = isneginf(x)
 
-    if not any(is_nan) and not any(is_pos_inf) and not any(is_neg_inf):
-        return
+    # if not any(is_nan) and not any(is_pos_inf) and not any(is_neg_inf):
+    #     return
 
     x = switch(is_nan, nan, x)
 
@@ -3140,6 +3156,8 @@ def nan_to_num(x, nan=0.0, posinf=None, neginf=None):
     "not_equal",
     "isnan",
     "isinf",
+    "isposinf",
+    "isneginf",
     "allclose",
     "isclose",
     "and_",

tests/tensor/test_math.py

@@ -79,6 +79,8 @@
     isinf,
     isnan,
     isnan_,
+    isneginf,
+    isposinf,
     log,
     log1mexp,
     log1p,
@@ -3644,6 +3646,36 @@ def test_grad_n_undefined(self):
             grad(polygamma(n, 0.5), wrt=n)
 
 
+def test_infs():
+    x = tensor(shape=(7,))
+
+    f_pos = function(
+        [x],
+        isposinf(x),
+        on_unused_input="warn",
+        allow_input_downcast=True,
+    )
+    f_neg = function(
+        [x],
+        isneginf(x),
+        on_unused_input="warn",
+        allow_input_downcast=True,
+    )
+
+    y = np.array([1, np.inf, 2, np.inf, -np.inf, -np.inf, 4])
+    out_pos = f_pos(y)
+    out_neg = f_neg(y)
+
+    np.testing.assert_allclose(
+        out_pos,
+        [0, 1, 0, 1, 0, 0, 0],
+    )
+    np.testing.assert_allclose(
+        out_neg,
+        [0, 0, 0, 0, 1, 1, 0],
+    )
+
+
 @pytest.mark.parametrize(
     ["nan", "posinf", "neginf"],
     [(0, None, None), (0, 0, 0), (0, None, 1000), (3, 1, -1)],
@@ -3655,7 +3687,7 @@ def test_nan_to_num(nan, posinf, neginf):
 
     f = function(
         [x],
-        nan_to_num(x, nan, posinf, neginf),
+        out,
         on_unused_input="warn",
         allow_input_downcast=True,
     )