Adding equivalent to numpy.nan_to_num functionality

[aspeers]

Adding equivalent to numpy.nan_to_num functionality
(https://numpy.org/doc/stable/reference/generated/numpy.nan_to_num.html)
addressing issue 479 (#479).

Motivation for these changes

Requested in issue #479

Implementation details

Added nan_to_num in pytensor/tensor/math.py

Checklist

• Explain motivation and implementation 👆
• Make sure that the pre-commit linting/style checks pass.
• Link relevant issues, preferably in nice commit messages.
• The commits correspond to relevant logical changes. Note that if they don't, we will rewrite/rebase/squash the git history before merging.
• Are the changes covered by tests and docstrings?
• Fill out the short summary sections 👇

Major / Breaking Changes

• ...

New features

• added nan_to_num

Bugfixes

• ...

Documentation

Example:

import pytensor
from pytensor import tensor as pt

import numpy as np

# Replace NaN
print("Replace NaN")
x = pt.dvector("x")
y = pt.log(x)

f = pytensor.function([x], y)

t = np.random.normal(size=(10,))
print("x")
print(t)
print("f = log(x)")
print(f (t))

# Test: NaN default replacement
print("\nTest: NaN default replacement")
y1 = pt.math.nan_to_num(y)
f = pytensor.function([x], y1)
print(f (t))

# Test: NaN custom replacement
print("\nTest: NaN custom replacement")
y2 = pt.math.nan_to_num(y, nan=1)
f = pytensor.function([x], y2)
print(f (t))

# Replace +INF/-INF
print("\n\nReplace +INF/-INF")
a = pt.dvector("a")
b = pt.dvector("b")
c = pt.math.true_div(a,b)

t_a = np.random.normal(size=(10,))
t_b = np.zeros((10,))

f = pytensor.function([a, b], c)
print("a")
print(t_a)
print("b")
print(t_b)
print("f = a/b")
print(f(t_a, t_b))

# Test: +INF/-INF default replacement
c1 = pt.math.nan_to_num(c)
f = pytensor.function([a, b], c1)
print("\nTest: +INF/-INF default replacement")
print(f(t_a, t_b))

# Test: +INF/-INF custom replacement
c2 = pt.math.nan_to_num(c, posinf=5, neginf=-5)
f = pytensor.function([a, b], c2)
print("\nTest: +INF/-INF custom replacement")
print(f(t_a, t_b))

Maintenance

• ...

[aspeers]

First PR from yesterday's code sprint. Main reasons for marking as in progress included:

• Had some difficulty running pytest locally. Resulted in multiple failures (most involving missing files/config parameters) which I figured the community could help by providing the relevant pytest command to run.
• Could use advice on adding unit tests for the addition.
• Was suggested to use ifelse() by @jessegrabowski instead of traditional if statements. ifelse is currently not imported/used in the math.py file and adding it to the imports list was causing a circular reference. Any thoughts would be welcomed!

[ricardoV94]

Thanks for opening the PR. Looks great!

I left some small technical suggestions. We will also need some tests

[ricardoV94]

Let's compare directly with posinf and neginf, if we don't have a isposinf and isneginf it's a good time to add it.

Suggested change

	x = switch(bitwise_and(isinf(x), pos), maxf, x)
	x = switch(bitwise_and(isinf(x), ~pos), minf, x)
	x = switch(bitwise_and(isinf(x), pos), maxf, x)
	x = switch(bitwise_and(isinf(x), ~pos), minf, x)

[ricardoV94]

Let's define it only when posinf and neginf are None

[ricardoV94]

Let's just copy Numpy docstring. No point in reinventing it

[ricardoV94]

Instead of overriding x, let's write all the checks isnan`isposinf\isneginf` on the original array.

[ricardoV94]

ifelse is not appropriate because it requires a scalar condition, and here we want to work with arbitrary tensors

[Dhruvanshu-Joshi]

Hi @aspeers , are you still working on this?