[OpenVINO backend] Support for nanarmax and nanar min #22383
+126
−8
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -340,6 +340,22 @@ def amin(x, axis=None, keepdims=False): | |
|
|
||
|
|
||
| def _resolve_axis(x, axis): | ||
| if hasattr(axis, "__class__") and "Output" in axis.__class__.__name__: | ||
| axis_node = axis.get_node() | ||
| if hasattr(axis_node, "get_data"): | ||
| axis_val = axis_node.get_data() | ||
| if axis_val is not None: | ||
| if ( | ||
| hasattr(axis_val, "ndim") | ||
| and axis_val.ndim > 0 | ||
| and axis_val.size > 0 | ||
| ): | ||
| axis = int(axis_val[0]) | ||
| elif hasattr(axis_val, "size") and axis_val.size == 1: | ||
| axis = int(axis_val.item()) | ||
| else: | ||
| axis = int(axis_val) | ||
|
|
||
|
Comment on lines
+343
to
+358
Collaborator
There was a problem hiding this comment. Why was this added? Tensor axes are not supported. |
||
| if axis == () or axis == []: | ||
| return x, None | ||
| if axis is None: | ||
|
|
@@ -522,6 +538,27 @@ def argmax(x, axis=None, keepdims=False): | |
| axis = 0 | ||
| k = ov_opset.constant(1, Type.i32).output(0) | ||
| else: | ||
| if isinstance(axis, ov.Output) or ( | ||
| hasattr(axis, "__class__") and "Output" in axis.__class__.__name__ | ||
| ): | ||
| axis_node = axis.get_node() | ||
| if hasattr(axis_node, "get_data"): | ||
| axis_val = axis_node.get_data() | ||
| if axis_val is not None: | ||
| if ( | ||
| hasattr(axis_val, "ndim") | ||
| and axis_val.ndim > 0 | ||
| and axis_val.size > 0 | ||
| ): | ||
| axis = int(axis_val[0]) | ||
| elif hasattr(axis_val, "size") and axis_val.size == 1: | ||
| axis = int(axis_val.item()) | ||
| else: | ||
| axis = int(axis_val) | ||
| else: | ||
| raise ValueError("axis must be static for argmax") | ||
|
Contributor
There was a problem hiding this comment. According to the Keras API design guidelines (line 140), error messages should be contextual and actionable. This error message could be improved by providing more context, such as the type of axis received and what is expected. A similar improvement can be applied to the For example: raise ValueError(
f"For the OpenVINO backend, `argmax` requires a static `axis` value, "
f"but received a dynamic tensor. Please provide a static integer "
f"for the `axis` argument."
)References
|
||
| else: | ||
| raise ValueError("axis must be static for argmax") | ||
|
Comment on lines
+541
to
+561
Collaborator
There was a problem hiding this comment. Why was this added? Tensor axes are not supported. |
||
| if axis < 0: | ||
| axis = rank + axis | ||
| k = ov_opset.constant(1, Type.i32).output(0) | ||
|
|
@@ -554,6 +591,27 @@ def argmin(x, axis=None, keepdims=False): | |
| axis = 0 | ||
| k = ov_opset.constant(1, Type.i32).output(0) | ||
| else: | ||
| if isinstance(axis, ov.Output) or ( | ||
| hasattr(axis, "__class__") and "Output" in axis.__class__.__name__ | ||
| ): | ||
| axis_node = axis.get_node() | ||
| if hasattr(axis_node, "get_data"): | ||
| axis_val = axis_node.get_data() | ||
| if axis_val is not None: | ||
| if ( | ||
| hasattr(axis_val, "ndim") | ||
| and axis_val.ndim > 0 | ||
| and axis_val.size > 0 | ||
| ): | ||
| axis = int(axis_val[0]) | ||
| elif hasattr(axis_val, "size") and axis_val.size == 1: | ||
| axis = int(axis_val.item()) | ||
| else: | ||
| axis = int(axis_val) | ||
| else: | ||
| raise ValueError("axis must be static for argmin") | ||
| else: | ||
| raise ValueError("axis must be static for argmin") | ||
|
Comment on lines
+594
to
+614
Collaborator
There was a problem hiding this comment. Why was this added? Tensor axes are not supported. |
||
| if axis < 0: | ||
| axis = rank + axis | ||
| k = ov_opset.constant(1, Type.i32).output(0) | ||
|
|
@@ -2848,15 +2906,79 @@ def moveaxis(x, source, destination): | |
|
|
||
|
|
||
| def nanargmax(x, axis=None, keepdims=False): | ||
| if isinstance(x, np.ndarray) and x.dtype == np.float64: | ||
| # conversion to f32 due to https://github.com/openvinotoolkit/openvino/issues/34138 | ||
| x = x.astype(np.float32) | ||
| x = get_ov_output(x) | ||
| x_type = x.get_element_type() | ||
| if x_type == Type.f64: | ||
| # conversion to f32 due to https://github.com/openvinotoolkit/openvino/issues/34138 | ||
| x = ov_opset.convert(x, Type.f32).output(0) | ||
| x_type = Type.f32 | ||
|
|
||
| if x_type.is_integral() or x_type == Type.boolean: | ||
| return argmax(OpenVINOKerasTensor(x), axis=axis, keepdims=keepdims) | ||
|
|
||
| x, axis = _resolve_axis(x, axis) | ||
| if axis is None: | ||
| return OpenVINOKerasTensor(x) | ||
|
Comment on lines
+2923
to
+2924
Contributor
There was a problem hiding this comment. When |
||
|
|
||
| nan_mask = ov_opset.is_nan(x) | ||
| neg_inf = ov_opset.constant(np.array(-np.inf, dtype=np.float32)) | ||
| if x_type != Type.f32: | ||
| neg_inf = ov_opset.convert(neg_inf, x_type) | ||
| x_replaced = ov_opset.select(nan_mask, neg_inf, x).output(0) | ||
|
|
||
| result = argmax( | ||
| OpenVINOKerasTensor(x_replaced), axis=axis, keepdims=keepdims | ||
| ) | ||
| result_ov = get_ov_output(result) | ||
|
|
||
| all_nan = ov_opset.reduce_logical_and(nan_mask, axis, keepdims).output(0) | ||
| nan_value = ov_opset.constant(-1, Type.i32).output(0) | ||
| if result_ov.get_element_type() != Type.i32: | ||
| nan_value = ov_opset.convert(nan_value, result_ov.get_element_type()) | ||
|
Comment on lines
+2939
to
+2940
Contributor
There was a problem hiding this comment. The |
||
| result_ov = ov_opset.select(all_nan, nan_value, result_ov).output(0) | ||
|
|
||
| return OpenVINOKerasTensor(result_ov) | ||
|
|
||
|
|
||
| def nanargmin(x, axis=None, keepdims=False): | ||
| if isinstance(x, np.ndarray) and x.dtype == np.float64: | ||
| # conversion to f32 due to https://github.com/openvinotoolkit/openvino/issues/34138 | ||
| x = x.astype(np.float32) | ||
| x = get_ov_output(x) | ||
| x_type = x.get_element_type() | ||
| if x_type == Type.f64: | ||
| # conversion to f32 due to https://github.com/openvinotoolkit/openvino/issues/34138 | ||
| x = ov_opset.convert(x, Type.f32).output(0) | ||
| x_type = Type.f32 | ||
|
|
||
| if x_type.is_integral() or x_type == Type.boolean: | ||
| return argmin(OpenVINOKerasTensor(x), axis=axis, keepdims=keepdims) | ||
|
|
||
| x, axis = _resolve_axis(x, axis) | ||
| if axis is None: | ||
| return OpenVINOKerasTensor(x) | ||
|
|
||
| nan_mask = ov_opset.is_nan(x) | ||
| pos_inf = ov_opset.constant(np.array(np.inf, dtype=np.float32)) | ||
| if x_type != Type.f32: | ||
| pos_inf = ov_opset.convert(pos_inf, x_type) | ||
| x_replaced = ov_opset.select(nan_mask, pos_inf, x).output(0) | ||
|
|
||
| result = argmin( | ||
| OpenVINOKerasTensor(x_replaced), axis=axis, keepdims=keepdims | ||
| ) | ||
| result_ov = get_ov_output(result) | ||
|
|
||
| all_nan = ov_opset.reduce_logical_and(nan_mask, axis, keepdims).output(0) | ||
| nan_value = ov_opset.constant(-1, Type.i32).output(0) | ||
| if result_ov.get_element_type() != Type.i32: | ||
| nan_value = ov_opset.convert(nan_value, result_ov.get_element_type()) | ||
| result_ov = ov_opset.select(all_nan, nan_value, result_ov).output(0) | ||
|
|
||
| return OpenVINOKerasTensor(result_ov) | ||
|
|
||
|
|
||
| def nancumsum(x, axis=None, dtype=None): | ||
|
|
||
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
This logic to resolve a static axis value from an OpenVINO tensor is duplicated in
argmax(lines 537-551) andargmin(lines 584-598). To improve maintainability and reduce code redundancy, consider extracting this logic into a private helper function. This function could attempt to resolve the axis to an integer and be used in all three places.Additionally, the type check
hasattr(axis, "__class__") and "Output" in axis.__class__.__name__is fragile. It's better to useisinstance(axis, ov.Output). Theargmaxandargminfunctions use a redundantorcondition with both checks. A single, robustisinstancecheck should be used in the new helper function.