[OpenVINO Backend] : add support for numpy.nan_to_num (#21186)

* feat: add support for numpy.nan_to_num

Signed-off-by: 11happy <[email protected]>

* use np.inf

Signed-off-by: 11happy <[email protected]>

* correct implementation based on new tests

Signed-off-by: 11happy <[email protected]>

* use np only torch having import errors

Signed-off-by: 11happy <[email protected]>

* use inf approach

Signed-off-by: 11happy <[email protected]>

* refactor code

Signed-off-by: 11happy <[email protected]>

---------

Signed-off-by: 11happy <[email protected]>

Author: 11happy

keras/src/backend/openvino/core.py

@@ -41,6 +41,38 @@
     "string": ov.Type.string,
 }
 
+DTYPES_MAX = {
+    ov.Type.bf16: 3.38953139e38,
+    ov.Type.f16: np.finfo(np.float16).max,
+    ov.Type.f32: np.finfo(np.float32).max,
+    ov.Type.f64: np.finfo(np.float64).max,
+    ov.Type.u8: np.iinfo(np.uint8).max,
+    ov.Type.u16: np.iinfo(np.uint16).max,
+    ov.Type.u32: np.iinfo(np.uint32).max,
+    ov.Type.u64: np.iinfo(np.uint64).max,
+    ov.Type.i8: np.iinfo(np.int8).max,
+    ov.Type.i16: np.iinfo(np.int16).max,
+    ov.Type.i32: np.iinfo(np.int32).max,
+    ov.Type.i64: np.iinfo(np.int64).max,
+    ov.Type.boolean: 1,
+}
+
+DTYPES_MIN = {
+    ov.Type.bf16: -3.38953139e38,
+    ov.Type.f16: np.finfo(np.float16).min,
+    ov.Type.f32: np.finfo(np.float32).min,
+    ov.Type.f64: np.finfo(np.float64).min,
+    ov.Type.u8: np.iinfo(np.uint8).min,
+    ov.Type.u16: np.iinfo(np.uint16).min,
+    ov.Type.u32: np.iinfo(np.uint32).min,
+    ov.Type.u64: np.iinfo(np.uint64).min,
+    ov.Type.i8: np.iinfo(np.int8).min,
+    ov.Type.i16: np.iinfo(np.int16).min,
+    ov.Type.i32: np.iinfo(np.int32).min,
+    ov.Type.i64: np.iinfo(np.int64).min,
+    ov.Type.boolean: 0,
+}
+
 
 def align_operand_types(x1, x2, op_name):
     x1_type = x1.element_type

keras/src/backend/openvino/excluded_concrete_tests.txt

@@ -38,7 +38,6 @@ NumpyDtypeTest::test_meshgrid
 NumpyDtypeTest::test_minimum_python_types
 NumpyDtypeTest::test_moveaxis
 NumpyDtypeTest::test_multiply
-NumpyDtypeTest::test_nan
 NumpyDtypeTest::test_outer_
 NumpyDtypeTest::test_power
 NumpyDtypeTest::test_prod
@@ -95,7 +94,6 @@ NumpyOneInputOpsCorrectnessTest::test_mean
 NumpyOneInputOpsCorrectnessTest::test_median
 NumpyOneInputOpsCorrectnessTest::test_meshgrid
 NumpyOneInputOpsCorrectnessTest::test_moveaxis
-NumpyOneInputOpsCorrectnessTest::test_nan_to_num
 NumpyOneInputOpsCorrectnessTest::test_pad_float16_constant_2
 NumpyOneInputOpsCorrectnessTest::test_pad_float32_constant_2
 NumpyOneInputOpsCorrectnessTest::test_pad_float64_constant_2

keras/src/backend/openvino/numpy.py

@@ -5,6 +5,8 @@
 from keras.src.backend import config
 from keras.src.backend.common import dtypes
 from keras.src.backend.common.variables import standardize_dtype
+from keras.src.backend.openvino.core import DTYPES_MAX
+from keras.src.backend.openvino.core import DTYPES_MIN
 from keras.src.backend.openvino.core import OPENVINO_DTYPES
 from keras.src.backend.openvino.core import OpenVINOKerasTensor
 from keras.src.backend.openvino.core import (
@@ -1079,9 +1081,36 @@ def moveaxis(x, source, destination):
 
 
 def nan_to_num(x, nan=0.0, posinf=None, neginf=None):
-    raise NotImplementedError(
-        "`nan_to_num` is not supported with openvino backend"
-    )
+    x = get_ov_output(x)
+    dtype = x.get_element_type()
+    if dtype.is_integral():
+        return OpenVINOKerasTensor(x)
+    isfloat64 = True if dtype == Type.f64 else False
+    if isfloat64:  # conversion to f32 due to https://github.com/openvinotoolkit/openvino/issues/30264
+        x = ov_opset.convert(x, Type.f32).output(0)
+        dtype = Type.f32
+    nan_val = ov_opset.constant(nan, dtype).output(0)
+    posinf_val = ov_opset.constant(
+        posinf if posinf is not None else DTYPES_MAX[dtype], dtype
+    ).output(0)
+    neginf_val = ov_opset.constant(
+        neginf if neginf is not None else DTYPES_MIN[dtype], dtype
+    ).output(0)
+    posinf_mask = ov_opset.is_inf(
+        x,
+        {"detect_positive": True, "detect_negative": False},
+    ).output(0)
+    neginf_mask = ov_opset.is_inf(
+        x,
+        {"detect_positive": False, "detect_negative": True},
+    ).output(0)
+    nan_mask = ov_opset.is_nan(x).output(0)
+    x = ov_opset.select(nan_mask, nan_val, x).output(0)
+    x = ov_opset.select(posinf_mask, posinf_val, x).output(0)
+    x = ov_opset.select(neginf_mask, neginf_val, x).output(0)
+    if isfloat64:
+        x = ov_opset.convert(x, Type.f64).output(0)
+    return OpenVINOKerasTensor(x)
 
 
 def ndim(x):