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,8 @@ 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)
-
-    if not any(is_nan) and not any(is_pos_inf) and not any(is_neg_inf):
-        return
+    is_pos_inf = isposinf(x)
+    is_neg_inf = isneginf(x)
 
     x = switch(is_nan, nan, x)
 
@@ -3140,6 +3153,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,26 @@ 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))
+    f_neg = function([x], isneginf(x))
+
+    y = np.array([1, np.inf, 2, np.inf, -np.inf, -np.inf, 4]).astype(x.dtype)
+    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)],
@@ -3653,14 +3675,9 @@ def test_nan_to_num(nan, posinf, neginf):
 
     out = nan_to_num(x, nan, posinf, neginf)
 
-    f = function(
-        [x],
-        nan_to_num(x, nan, posinf, neginf),
-        on_unused_input="warn",
-        allow_input_downcast=True,
-    )
+    f = function([x], out)
 
-    y = np.array([1, 2, np.nan, np.inf, -np.inf, 3, 4])
+    y = np.array([1, 2, np.nan, np.inf, -np.inf, 3, 4]).astype(x.dtype)
     out = f(y)
 
     posinf = np.finfo(x.real.dtype).max if posinf is None else posinf