From 829309bcccdf5cc8b67c76f2230d4dfd3c3fda3c Mon Sep 17 00:00:00 2001 From: Dhruvanshu-Joshi Date: Mon, 6 May 2024 08:32:16 +0530 Subject: [PATCH 1/3] Add nan_to_num conversion --- pytensor/scalar/basic.py | 50 +++++++++++++++++++ pytensor/tensor/math.py | 100 ++++++++++++++++++++++++++++++++++++++ tests/tensor/test_math.py | 29 +++++++++++ 3 files changed, 179 insertions(+) diff --git a/pytensor/scalar/basic.py b/pytensor/scalar/basic.py index 56a3629dc5..9fc868c31c 100644 --- a/pytensor/scalar/basic.py +++ b/pytensor/scalar/basic.py @@ -1533,6 +1533,56 @@ def c_code_cache_version(self): isinf = IsInf() +class IsPosInf(FixedLogicalComparison): + nfunc_spec = ("isposinf", 1, 1) + + def impl(self, x): + return np.isposinf(x) + + def c_code(self, node, name, inputs, outputs, sub): + (x,) = inputs + (z,) = outputs + if node.inputs[0].type in complex_types: + raise NotImplementedError() + # Discrete type can never be posinf + if node.inputs[0].type in discrete_types: + return f"{z} = false;" + + return f"{z} = isinf({x}) && !signbit({x});" + + def c_code_cache_version(self): + scalarop_version = super().c_code_cache_version() + return (*scalarop_version, 4) + + +isposinf = IsPosInf() + + +class IsNegInf(FixedLogicalComparison): + nfunc_spec = ("isneginf", 1, 1) + + def impl(self, x): + return np.isneginf(x) + + def c_code(self, node, name, inputs, outputs, sub): + (x,) = inputs + (z,) = outputs + if node.inputs[0].type in complex_types: + raise NotImplementedError() + # Discrete type can never be neginf + if node.inputs[0].type in discrete_types: + return f"{z} = false;" + + return f"{z} = isinf({x}) && signbit({x});" + + def c_code_cache_version(self): + scalarop_version = super().c_code_cache_version() + return (*scalarop_version, 4) + + +isneginf = IsNegInf() + + class InRange(LogicalComparison): nin = 3 diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py index 63a943e1f1..fbf4cd1e93 100644 --- a/pytensor/tensor/math.py +++ b/pytensor/tensor/math.py @@ -881,6 +881,46 @@ def isinf(a): return isinf_(a) +@scalar_elemwise +def isposinf(a): + """isposinf(a)""" + + +# Rename isposnan to isposnan_ to allow to bypass it when not needed. +# glibc 2.23 don't allow isposnan on int, so we remove it from the graph. +isposinf_ = isposinf + + +def isposinf(a): + """isposinf(a)""" + a = as_tensor_variable(a) + if a.dtype in discrete_dtypes: + return alloc( + np.asarray(False, dtype="bool"), *[a.shape[i] for i in range(a.ndim)] + ) + return isposinf_(a) + + +@scalar_elemwise +def isneginf(a): + """isneginf(a)""" + + +# Rename isnegnan to isnegnan_ to allow to bypass it when not needed. +# glibc 2.23 don't allow isnegnan on int, so we remove it from the graph. +isneginf_ = isneginf + + +def isneginf(a): + """isneginf(a)""" + a = as_tensor_variable(a) + if a.dtype in discrete_dtypes: + return alloc( + np.asarray(False, dtype="bool"), *[a.shape[i] for i in range(a.ndim)] + ) + return isneginf_(a) + + def allclose(a, b, rtol=1.0e-5, atol=1.0e-8, equal_nan=False): """ Implement Numpy's ``allclose`` on tensors. @@ -3043,6 +3083,65 @@ def vectorize_node_dot_to_matmul(op, node, batched_x, batched_y): return vectorize_node_fallback(op, node, batched_x, batched_y) +def nan_to_num(x, nan=0.0, posinf=None, neginf=None): + """ + Replace NaN with zero and infinity with large finite numbers (default + behaviour) or with the numbers defined by the user using the `nan`, + `posinf` and/or `neginf` keywords. + + NaN is replaced by zero or by the user defined value in + `nan` keyword, infinity is replaced by the largest finite floating point + values representable by ``x.dtype`` or by the user defined value in + `posinf` keyword and -infinity is replaced by the most negative finite + floating point values representable by ``x.dtype`` or by the user defined + value in `neginf` keyword. + + Parameters + ---------- + x : symbolic tensor + Input array. + nan + The value to replace NaN's with in the tensor (default = 0). + posinf + The value to replace +INF with in the tensor (default max + in range representable by ``x.dtype``). + neginf + The value to replace -INF with in the tensor (default min + in range representable by ``x.dtype``). + + Returns + ------- + out + The tensor with NaN's, +INF, and -INF replaced with the + specified and/or default substitutions. + """ + # Replace NaN's with nan keyword + is_nan = isnan(x) + 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 + + x = switch(is_nan, nan, x) + + # Get max and min values representable by x.dtype + maxf = posinf + minf = neginf + + # Specify the value to replace +INF and -INF with + if maxf is None: + maxf = np.finfo(x.real.dtype).max + if minf is None: + minf = np.finfo(x.real.dtype).min + + # Replace +INF and -INF values + x = switch(is_pos_inf, maxf, x) + x = switch(is_neg_inf, minf, x) + + return x + + # NumPy logical aliases square = sqr @@ -3199,4 +3298,5 @@ def vectorize_node_dot_to_matmul(op, node, batched_x, batched_y): "logaddexp", "logsumexp", "hyp2f1", + "nan_to_num", ] diff --git a/tests/tensor/test_math.py b/tests/tensor/test_math.py index e346348406..9dd02ea68c 100644 --- a/tests/tensor/test_math.py +++ b/tests/tensor/test_math.py @@ -95,6 +95,7 @@ minimum, mod, mul, + nan_to_num, neg, neq, outer, @@ -3641,3 +3642,31 @@ def test_grad_n_undefined(self): n = scalar(dtype="int64") with pytest.raises(NullTypeGradError): grad(polygamma(n, 0.5), wrt=n) + + +@pytest.mark.parametrize( + ["nan", "posinf", "neginf"], + [(0, None, None), (0, 0, 0), (0, None, 1000), (3, 1, -1)], +) +def test_nan_to_num(nan, posinf, neginf): + x = tensor(shape=(7,)) + + 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, + ) + + y = np.array([1, 2, np.nan, np.inf, -np.inf, 3, 4]) + out = f(y) + + posinf = np.finfo(x.real.dtype).max if posinf is None else posinf + neginf = np.finfo(x.real.dtype).min if neginf is None else neginf + + np.testing.assert_allclose( + out, + np.nan_to_num(y, nan=nan, posinf=posinf, neginf=neginf), + ) From 64e60edc67016aa066506820736e88f85244289a Mon Sep 17 00:00:00 2001 From: Dhruvanshu-Joshi Date: Sun, 9 Jun 2024 14:08:11 +0530 Subject: [PATCH 2/3] Replaced use of isposinf and isneginf op with existing ops --- pytensor/scalar/basic.py | 50 ---------------------------------------- pytensor/tensor/math.py | 44 ++--------------------------------- 2 files changed, 2 insertions(+), 92 deletions(-) diff --git a/pytensor/scalar/basic.py b/pytensor/scalar/basic.py index 9fc868c31c..56a3629dc5 100644 --- a/pytensor/scalar/basic.py +++ b/pytensor/scalar/basic.py @@ -1533,56 +1533,6 @@ def c_code_cache_version(self): isinf = IsInf() -class IsPosInf(FixedLogicalComparison): - nfunc_spec = ("isposinf", 1, 1) - - def impl(self, x): - return np.isposinf(x) - - def c_code(self, node, name, inputs, outputs, sub): - (x,) = inputs - (z,) = outputs - if node.inputs[0].type in complex_types: - raise NotImplementedError() - # Discrete type can never be posinf - if node.inputs[0].type in discrete_types: - return f"{z} = false;" - - return f"{z} = isinf({x}) && !signbit({x});" - - def c_code_cache_version(self): - scalarop_version = super().c_code_cache_version() - return (*scalarop_version, 4) - - -isposinf = IsPosInf() - - -class IsNegInf(FixedLogicalComparison): - nfunc_spec = ("isneginf", 1, 1) - - def impl(self, x): - return np.isneginf(x) - - def c_code(self, node, name, inputs, outputs, sub): - (x,) = inputs - (z,) = outputs - if node.inputs[0].type in complex_types: - raise NotImplementedError() - # Discrete type can never be neginf - if node.inputs[0].type in discrete_types: - return f"{z} = false;" - - return f"{z} = isinf({x}) && signbit({x});" - - def c_code_cache_version(self): - scalarop_version = super().c_code_cache_version() - return (*scalarop_version, 4) - - -isneginf = IsNegInf() - - class InRange(LogicalComparison): nin = 3 diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py index fbf4cd1e93..1feac6bd15 100644 --- a/pytensor/tensor/math.py +++ b/pytensor/tensor/math.py @@ -881,46 +881,6 @@ def isinf(a): return isinf_(a) -@scalar_elemwise -def isposinf(a): - """isposinf(a)""" - - -# Rename isposnan to isposnan_ to allow to bypass it when not needed. -# glibc 2.23 don't allow isposnan on int, so we remove it from the graph. -isposinf_ = isposinf - - -def isposinf(a): - """isposinf(a)""" - a = as_tensor_variable(a) - if a.dtype in discrete_dtypes: - return alloc( - np.asarray(False, dtype="bool"), *[a.shape[i] for i in range(a.ndim)] - ) - return isposinf_(a) - - -@scalar_elemwise -def isneginf(a): - """isneginf(a)""" - - -# Rename isnegnan to isnegnan_ to allow to bypass it when not needed. -# glibc 2.23 don't allow isnegnan on int, so we remove it from the graph. -isneginf_ = isneginf - - -def isneginf(a): - """isneginf(a)""" - a = as_tensor_variable(a) - if a.dtype in discrete_dtypes: - return alloc( - np.asarray(False, dtype="bool"), *[a.shape[i] for i in range(a.ndim)] - ) - return isneginf_(a) - - def allclose(a, b, rtol=1.0e-5, atol=1.0e-8, equal_nan=False): """ Implement Numpy's ``allclose`` on tensors. @@ -3117,8 +3077,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 = isposinf(x) - is_neg_inf = isneginf(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 From 7c68b299ee6cfd536ee1efdf6481fe10ae9adb5c Mon Sep 17 00:00:00 2001 From: Dhruvanshu-Joshi Date: Mon, 17 Jun 2024 16:46:05 +0530 Subject: [PATCH 3/3] Added numpy-like posinf and neginf --- pytensor/tensor/math.py | 25 ++++++++++++++++++++----- tests/tensor/test_math.py | 31 ++++++++++++++++++++++++------- 2 files changed, 44 insertions(+), 12 deletions(-) diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py index 1feac6bd15..132642e70e 100644 --- a/pytensor/tensor/math.py +++ b/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_", diff --git a/tests/tensor/test_math.py b/tests/tensor/test_math.py index 9dd02ea68c..e1f9465d1e 100644 --- a/tests/tensor/test_math.py +++ b/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