Skip to content

MNT helper _nan_to_num function for array API #30637

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 17 additions & 0 deletions sklearn/utils/_array_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -1108,3 +1108,20 @@
return array.tolist()
array_np = _convert_to_numpy(array, xp=xp)
return [element.item() for element in array_np]


def _nan_to_num(array, xp=None):
"""Substitutes NaN values of an array with 0 and inf values with the maximum or
minimum numbers available for the dtype respectively; like np.nan_to_num."""
xp, _ = get_namespace(array, xp=xp)
try:
array = xp.nan_to_num(array)
except AttributeError: # currently catching exceptions from array_api_strict
array[xp.isnan(array)] = 0
if xp.isdtype(array.dtype, "real floating"):
array[xp.isinf(array) & (array > 0)] = xp.finfo(array.dtype).max
array[xp.isinf(array) & (array < 0)] = xp.finfo(array.dtype).min

Check warning on line 1123 in sklearn/utils/_array_api.py

View check run for this annotation

Codecov / codecov/patch

sklearn/utils/_array_api.py#L1116-L1123

Added lines #L1116 - L1123 were not covered by tests
else: # xp.isdtype(array.dtype, "integral")
array[xp.isinf(array) & (array > 0)] = xp.iinfo(array.dtype).max
array[xp.isinf(array) & (array < 0)] = xp.iinfo(array.dtype).min
return array

Check warning on line 1127 in sklearn/utils/_array_api.py

View check run for this annotation

Codecov / codecov/patch

sklearn/utils/_array_api.py#L1125-L1127

Added lines #L1125 - L1127 were not covered by tests
Loading