Correct implementation for several OpenVINO operations (#21752)

* fix absolute, enable all/any

* fix ceil func

* enable input/output tests for any/all, fix axis resolutoin

* fix sqrt for int32

* fix sqrt and sum, consolidate duplicated logic

* fix squeeze

* add pos/neg/regular inf

* fix nan, add finite support

* add a note about why we're not using openvino's existing capabilities for finite/infinite

Author: danielenricocahall

keras/src/backend/openvino/excluded_concrete_tests.txt

@@ -1,9 +1,6 @@
 NumPyTestRot90
-NumpyDtypeTest::test_absolute_bool
 NumpyDtypeTest::test_add_
-NumpyDtypeTest::test_all
 NumpyDtypeTest::test_angle
-NumpyDtypeTest::test_any
 NumpyDtypeTest::test_argpartition
 NumpyDtypeTest::test_array
 NumpyDtypeTest::test_bartlett
@@ -16,7 +13,6 @@ NumpyDtypeTest::test_hypot
 NumpyDtypeTest::test_kaiser
 NumpyDtypeTest::test_bitwise
 NumpyDtypeTest::test_cbrt
-NumpyDtypeTest::test_ceil
 NumpyDtypeTest::test_concatenate
 NumpyDtypeTest::test_corrcoef
 NumpyDtypeTest::test_correlate
@@ -29,11 +25,7 @@ NumpyDtypeTest::test_exp2
 NumpyDtypeTest::test_flip
 NumpyDtypeTest::test_floor_divide
 NumpyDtypeTest::test_inner
-NumpyDtypeTest::test_isfinite
 NumpyDtypeTest::test_isin
-NumpyDtypeTest::test_isinf
-NumpyDtypeTest::test_isnan
-NumpyDtypeTest::test_isposinf
 NumpyDtypeTest::test_isreal
 NumpyDtypeTest::test_kron
 NumpyDtypeTest::test_lcm
@@ -48,10 +40,8 @@ NumpyDtypeTest::test_roll
 NumpyDtypeTest::test_round
 NumpyDtypeTest::test_searchsorted
 NumpyDtypeTest::test_signbit
-NumpyDtypeTest::test_sqrt
 NumpyDtypeTest::test_std
 NumpyDtypeTest::test_subtract
-NumpyDtypeTest::test_sum
 NumpyDtypeTest::test_swapaxes
 NumpyDtypeTest::test_tensordot_
 NumpyDtypeTest::test_tile
@@ -61,12 +51,8 @@ NumpyDtypeTest::test_unravel
 NumpyDtypeTest::test_var
 NumpyDtypeTest::test_vdot
 NumpyDtypeTest::test_vstack
-NumpyDtypeTest::test_clip_bool
-NumpyDtypeTest::test_square_bool
 HistogramTest
-NumpyOneInputOpsCorrectnessTest::test_all
 NumpyOneInputOpsCorrectnessTest::test_angle
-NumpyOneInputOpsCorrectnessTest::test_any
 NumpyOneInputOpsCorrectnessTest::test_argpartition
 NumpyOneInputOpsCorrectnessTest::test_array
 NumpyOneInputOpsCorrectnessTest::test_bartlett
@@ -86,9 +72,6 @@ NumpyOneInputOpsCorrectnessTest::test_exp2
 NumpyOneInputOpsCorrectnessTest::test_flip
 NumpyOneInputOpsCorrectnessTest::test_floor_divide
 NumpyOneInputOpsCorrectnessTest::test_imag
-NumpyOneInputOpsCorrectnessTest::test_isfinite
-NumpyOneInputOpsCorrectnessTest::test_isinf
-NumpyOneInputOpsCorrectnessTest::test_isposinf
 NumpyOneInputOpsCorrectnessTest::test_isreal
 NumpyOneInputOpsCorrectnessTest::test_logaddexp2
 NumpyOneInputOpsCorrectnessTest::test_pad_float16_constant_2
@@ -107,10 +90,7 @@ NumpyOneInputOpsCorrectnessTest::test_select
 NumpyOneInputOpsCorrectnessTest::test_signbit
 NumpyOneInputOpsCorrectnessTest::test_size
 NumpyOneInputOpsCorrectnessTest::test_slogdet
-NumpyOneInputOpsCorrectnessTest::test_sqrt_int32
-NumpyOneInputOpsCorrectnessTest::test_squeeze
 NumpyOneInputOpsCorrectnessTest::test_std
-NumpyOneInputOpsCorrectnessTest::test_sum
 NumpyOneInputOpsCorrectnessTest::test_swapaxes
 NumpyOneInputOpsCorrectnessTest::test_tile
 NumpyOneInputOpsCorrectnessTest::test_trace

keras/src/backend/openvino/numpy.py

@@ -181,6 +181,9 @@ def zeros(shape, dtype=None):
 
 def absolute(x):
     x = get_ov_output(x)
+    x_type = x.get_element_type()
+    if x_type == Type.boolean:
+        return OpenVINOKerasTensor(x)
     return OpenVINOKerasTensor(ov_opset.absolute(x).output(0))
 
 
@@ -191,11 +194,10 @@ def abs(x):
 
 def all(x, axis=None, keepdims=False):
     x = get_ov_output(x)
+    x, axis = _resolve_axis(x, axis)
     if axis is None:
-        flatten_shape = ov_opset.constant([-1], Type.i32).output(0)
-        x = ov_opset.reshape(x, flatten_shape, False).output(0)
-        axis = 0
-    axis = ov_opset.constant(axis, Type.i32).output(0)
+        return OpenVINOKerasTensor(x)
+    x = ov_opset.convert(x, Type.boolean).output(0)
     return OpenVINOKerasTensor(
         ov_opset.reduce_logical_and(x, axis, keepdims).output(0)
     )
@@ -207,11 +209,10 @@ def angle(x):
 
 def any(x, axis=None, keepdims=False):
     x = get_ov_output(x)
+    x, axis = _resolve_axis(x, axis)
     if axis is None:
-        flatten_shape = ov_opset.constant([-1], Type.i32).output(0)
-        x = ov_opset.reshape(x, flatten_shape, False).output(0)
-        axis = 0
-    axis = ov_opset.constant(axis, Type.i32).output(0)
+        return OpenVINOKerasTensor(x)
+    x = ov_opset.convert(x, Type.boolean).output(0)
     return OpenVINOKerasTensor(
         ov_opset.reduce_logical_or(x, axis, keepdims).output(0)
     )
@@ -256,6 +257,17 @@ def _resolve_axis(x, axis):
     return x, axis
 
 
+def _upcast_type_if_needed(x):
+    x_type = x.get_element_type()
+    if x_type == Type.boolean:
+        x = ov_opset.convert(x, Type.i32).output(0)
+    elif x_type in (Type.i8, Type.i16):
+        x = ov_opset.convert(x, Type.i32).output(0)
+    elif x_type in (Type.u8, Type.u16):
+        x = ov_opset.convert(x, Type.u32).output(0)
+    return x
+
+
 def append(x1, x2, axis=None):
     x1, x2 = get_ov_output(x1), get_ov_output(x2)
     x1, x2 = _align_operand_types(x1, x2, "append()")
@@ -616,11 +628,18 @@ def cbrt(x):
 
 def ceil(x):
     x = get_ov_output(x)
-    return OpenVINOKerasTensor(ov_opset.ceil(x).output(0))
+    x_type = x.get_element_type()
+    if x_type.is_integral():
+        x = ov_opset.convert(x, OPENVINO_DTYPES[config.floatx()]).output(0)
+    ceiling = ov_opset.ceil(x).output(0)
+    return OpenVINOKerasTensor(ceiling)
 
 
 def clip(x, x_min, x_max):
     x = get_ov_output(x)
+    x_type = x.get_element_type()
+    if x_type == Type.boolean:
+        x = ov_opset.convert(x, Type.i32).output(0)
     x_min = get_ov_output(x_min, x.get_element_type())
     x_max = get_ov_output(x_max, x.get_element_type())
     clip_by_min = ov_opset.maximum(x, x_min).output(0)
@@ -1012,25 +1031,52 @@ def isclose(x1, x2, rtol=1e-5, atol=1e-8, equal_nan=False):
 
 
 def isfinite(x):
-    x = get_ov_output(x)
-    return OpenVINOKerasTensor(ov_opset.is_finite(x).output(0))
+    # NOTE: openvino has an is_finite operation but it does not properly
+    # catch np.inf and -np.inf as not finite values. Hence we bootstrap here. If
+    # that ever changes, we could simplify this to just call that operation.
+    inf_values = get_ov_output(isinf(x))
+    nan_values = get_ov_output(isnan(x))
+    all_non_finite_values = ov_opset.logical_or(inf_values, nan_values).output(
+        0
+    )
+    is_finite = ov_opset.logical_not(all_non_finite_values).output(0)
+    return OpenVINOKerasTensor(is_finite)
 
 
 def isin(x1, x2, assume_unique=False, invert=False):
     raise NotImplementedError("`isin` is not supported with openvino backend")
 
 
 def isinf(x):
-    x = get_ov_output(x)
-    return OpenVINOKerasTensor(ov_opset.is_inf(x).output(0))
+    pos_inf = get_ov_output(isposinf(x))
+    neg_inf = get_ov_output(isneginf(x))
+    inf = ov_opset.logical_or(pos_inf, neg_inf).output(0)
+    return OpenVINOKerasTensor(inf)
 
 
 def isnan(x):
     x = get_ov_output(x)
+    x_type = x.get_element_type()
+    if x_type.is_integral():
+        x = ov_opset.convert(x, OPENVINO_DTYPES[config.floatx()])
     return OpenVINOKerasTensor(ov_opset.is_nan(x).output(0))
 
 
 def isneginf(x):
+    return _is_inf(x, pos=False)
+
+
+def isposinf(x):
+    return _is_inf(x)
+
+
+def _is_inf(x, pos=True):
+    # NOTE: there is an ov_opset.is_inf but it does not catch
+    # numpy infinite values like np.inf and -np.inf, hence why we have this
+    # if this ever changes in OpenVINO, we can do this instead:
+    # ov_opset.is_inf(x, {"detect_positive": pos, "detect_negative": not pos})
+    # for each infinite sign
+    inf_value = np.inf if pos else -np.inf
     x = get_ov_output(x)
     x_type = x.get_element_type()
 
@@ -1043,26 +1089,19 @@ def isneginf(x):
 
     if x_type == Type.bf16:
         x_f32 = ov_opset.convert(x, Type.f32).output(0)
-        neg_inf = ov_opset.constant(-np.inf, Type.f32).output(0)
-        is_neg_inf = ov_opset.equal(x_f32, neg_inf).output(0)
+        inf = ov_opset.constant(inf_value, Type.f32).output(0)
+        is_inf = ov_opset.equal(x_f32, inf).output(0)
     else:
         if x_type == Type.f16:
-            neg_inf = ov_opset.constant(-np.inf, Type.f16).output(0)
+            inf = ov_opset.constant(inf_value, Type.f16).output(0)
         elif x_type == Type.f32:
-            neg_inf = ov_opset.constant(-np.inf, Type.f32).output(0)
+            inf = ov_opset.constant(inf_value, Type.f32).output(0)
         elif x_type == Type.f64:
-            neg_inf = ov_opset.constant(-np.inf, Type.f64).output(0)
+            inf = ov_opset.constant(inf_value, Type.f64).output(0)
         else:
-            neg_inf = ov_opset.constant(-np.inf, Type.f32).output(0)
-        is_neg_inf = ov_opset.equal(x, neg_inf).output(0)
-
-    return OpenVINOKerasTensor(is_neg_inf)
-
-
-def isposinf(x):
-    raise NotImplementedError(
-        "`isposinf` is not supported with openvino backend"
-    )
+            inf = ov_opset.constant(inf_value, Type.f32).output(0)
+        is_inf = ov_opset.equal(x, inf).output(0)
+    return OpenVINOKerasTensor(is_inf)
 
 
 def isreal(x):
@@ -1746,23 +1785,10 @@ def prod(x, axis=None, keepdims=False, dtype=None):
         x = ov_opset.convert(x, ov_dtype).output(0)
     # Otherwise, apply dtype promotion rules before reduction.
     else:
-        x_type = x.get_element_type()
-        if x_type == Type.boolean:
-            x = ov_opset.convert(x, Type.i32).output(0)
-        elif x_type in (Type.i8, Type.i16):
-            x = ov_opset.convert(x, Type.i32).output(0)
-        elif x_type in (Type.u8, Type.u16):
-            x = ov_opset.convert(x, Type.u32).output(0)
-
+        x = _upcast_type_if_needed(x)
+    x, axis = _resolve_axis(x, axis)
     if axis is None:
-        flatten_shape = ov_opset.constant([-1], Type.i32).output(0)
-        x = ov_opset.reshape(x, flatten_shape, False).output(0)
-        axis = 0
-
-    if isinstance(axis, tuple):
-        axis = list(axis)
-    axis = ov_opset.constant(axis, Type.i32).output(0)
-
+        return OpenVINOKerasTensor(x)
     # Compute the product
     result = ov_opset.reduce_prod(x, axis, keepdims).output(0)
 
@@ -2323,12 +2349,19 @@ def negative(x):
 
 def square(x):
     x = get_ov_output(x)
+    x_type = x.get_element_type()
+    if x_type == Type.boolean:
+        x = ov_opset.convert(x, Type.i32).output(0)
     const_two = ov_opset.constant(2, x.get_element_type()).output(0)
     return OpenVINOKerasTensor(ov_opset.power(x, const_two).output(0))
 
 
 def sqrt(x):
     x = get_ov_output(x)
+    x_type = x.get_element_type()
+    if x_type.is_integral():
+        ov_type = OPENVINO_DTYPES[config.floatx()]
+        x = ov_opset.convert(x, ov_type).output(0)
     return OpenVINOKerasTensor(ov_opset.sqrt(x).output(0))
 
 
@@ -2339,6 +2372,8 @@ def squeeze(x, axis=None):
         for idx, dim in enumerate(x.get_partial_shape()):
             if dim == 1:
                 axis.append(idx)
+    if isinstance(axis, tuple):
+        axis = list(axis)
     axis = ov_opset.constant(axis, Type.i32).output(0)
     return OpenVINOKerasTensor(ov_opset.squeeze(x, axis).output(0))
 
@@ -2385,12 +2420,12 @@ def var(x, axis=None, keepdims=False):
 
 def sum(x, axis=None, keepdims=False):
     x = get_ov_output(x)
+    x, axis = _resolve_axis(x, axis)
     if axis is None:
-        flatten_shape = ov_opset.constant([-1], Type.i32).output(0)
-        x = ov_opset.reshape(x, flatten_shape, False).output(0)
-        axis = 0
-    axis = ov_opset.constant(axis, Type.i32).output(0)
-    return OpenVINOKerasTensor(ov_opset.reduce_sum(x, axis, keepdims).output(0))
+        return OpenVINOKerasTensor(x)
+    x = _upcast_type_if_needed(x)
+    summed_value = ov_opset.reduce_sum(x, axis, keepdims).output(0)
+    return OpenVINOKerasTensor(summed_value)
 
 
 def eye(N, M=None, k=0, dtype=None):