[OpenVINO backend] Support numpy.diagonal issue 29115 #21584
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). View this failed invocation of the CLA check for more information. For the most up to date status, view the checks section at the bottom of the pull request. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @arjunverma2004, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request introduces support for the numpy.diagonal operation within the OpenVINO backend for Keras 3. This enhancement allows Keras models leveraging the OpenVINO backend to correctly process operations involving extracting diagonals from arrays, thereby expanding the compatibility and functionality of the backend.
Highlights
- OpenVINO Backend Support: Implemented support for numpy.diagonal in the Keras 3 OpenVINO backend.
- diagonal Operation Decomposition: Added a detailed decomposition for the diagonal operation in keras/src/backend/openvino/numpy.py, translating the NumPy operation into a series of OpenVINO opset operations.
- Test Enablement: Removed test_diagonal from the excluded_concrete_tests.txt file, enabling the corresponding correctness tests for the diagonal operation.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Code Review
This pull request adds support for numpy.diagonal to the OpenVINO backend. The implementation is a decomposition of the operation using OpenVINO ops, which is the correct approach for this backend. My review focuses on improving code consistency and cleanliness by addressing issues with imports, unused code, and inconsistent operator usage. I've provided a single comprehensive suggestion to address these points.
keras/src/backend/openvino/numpy.py
Outdated
| # --- Chnage for issue 29115 --- | ||
| import openvino.runtime.opset14 as ov | ||
|
|
||
| from .core import OpenVINOKerasTensor # already present in file | ||
| from .core import _convert_to_node, _wrap_node # adapt if your file names differ | ||
|
|
||
| def diagonal(x, offset=0, axis1=0, axis2=1): | ||
| """OpenVINO backend decomposition for keras.ops.diagonal.""" | ||
| x_node = _convert_to_node(x) # -> ov.Node | ||
| offset_const = ov.constant(int(offset), dtype="i64") | ||
|
|
||
| # rank & normalize axes | ||
| shape = ov.shape_of(x_node) # i64 vector | ||
| rank = ov.shape_of(shape) # scalar i64 (len of shape) | ||
| rank_val = ov.squeeze(rank) # [] -> scalar | ||
| axis1_node = ov.mod(ov.add(ov.constant(int(axis1), dtype="i64"), rank_val), rank_val) | ||
| axis2_node = ov.mod(ov.add(ov.constant(int(axis2), dtype="i64"), rank_val), rank_val) | ||
|
|
||
| # If axis1 == axis2, behavior should match numpy error; Keras tests don't hit this, | ||
| # so we skip explicit assert to keep graph-friendly. | ||
|
|
||
| # Build permutation to move axis1, axis2 to the end | ||
| # perm = [all axes except axis1/axis2 in order] + [axis1, axis2] | ||
| arange = ov.range(ov.constant(0, dtype="i64"), rank_val, ov.constant(1, dtype="i64")) | ||
| mask1 = ov.equal(arange, axis1_node) | ||
| mask2 = ov.equal(arange, axis2_node) | ||
| not12 = ov.logical_not(ov.logical_or(mask1, mask2)) | ||
| others = ov.squeeze(ov.non_zero(not12), [1]) # gather positions != axis1, axis2 | ||
| perm = ov.concat([others, ov.reshape(axis1_node, [1]), ov.reshape(axis2_node, [1])], 0) | ||
|
|
||
| x_perm = ov.transpose(x_node, perm) | ||
| permuted_shape = ov.shape_of(x_perm) | ||
| # last two dims | ||
| last2 = ov.gather(permuted_shape, ov.constant([-2, -1], dtype="i64"), ov.constant(0, dtype="i64")) | ||
| d1 = ov.gather(permuted_shape, ov.constant([-2], dtype="i64"), ov.constant(0, dtype="i64")) | ||
| d2 = ov.gather(permuted_shape, ov.constant([-1], dtype="i64"), ov.constant(0, dtype="i64")) | ||
| d1 = ov.squeeze(d1) # scalar | ||
| d2 = ov.squeeze(d2) # scalar | ||
|
|
||
| # start1 = max(0, offset), start2 = max(0, -offset) | ||
| zero = ov.constant(0, dtype="i64") | ||
| start1 = ov.maximum(zero, offset_const) | ||
| start2 = ov.maximum(zero, ov.negative(offset_const)) | ||
|
|
||
| # L = min(d1 - start1, d2 - start2) | ||
| l1 = ov.subtract(d1, start1) | ||
| l2 = ov.subtract(d2, start2) | ||
| L = ov.minimum(l1, l2) | ||
|
|
||
| # r = range(0, L, 1) -> shape [L] | ||
| r = ov.range(zero, L, ov.constant(1, dtype="i64")) | ||
| idx_row = ov.add(r, start1) | ||
| idx_col = ov.add(r, start2) | ||
| idx_row = ov.unsqueeze(idx_row, ov.constant(1, dtype="i64")) # [L,1] | ||
| idx_col = ov.unsqueeze(idx_col, ov.constant(1, dtype="i64")) # [L,1] | ||
| diag_idx = ov.concat([idx_row, idx_col], 1) # [L,2] | ||
|
|
||
| # Broadcast indices to batch dims: target shape = (*batch, L, 2) | ||
| # batch_rank = rank(x) - 2 | ||
| two = ov.constant(2, dtype="i64") | ||
| batch_rank = ov.subtract(rank_val, two) | ||
| # build target shape: concat(permuted_shape[:batch_rank], [L, 2]) | ||
| batch_shape = ov.slice(permuted_shape, ov.constant([0], dtype="i64"), | ||
| ov.reshape(batch_rank, [1]), ov.constant([1], dtype="i64")) | ||
| target_shape = ov.concat([batch_shape, ov.reshape(L, [1]), ov.constant([2], dtype="i64")], 0) | ||
| bcast_idx = ov.broadcast(diag_idx, target_shape) | ||
|
|
||
| # GatherND with batch_dims = batch_rank | ||
| gathered = ov.gather_nd(x_perm, bcast_idx, batch_rank) | ||
|
|
||
| return OpenVINOKerasTensor(gathered) |
There was a problem hiding this comment.
The new diagonal function and its surrounding code have a few areas for improvement regarding consistency and code cleanliness:
- Imports and Comments: The introductory comment has a typo and is not needed. There are also redundant/unused imports and developer-facing comments that should be removed.
- Inconsistent API Usage: A new import alias
ovis used, while the rest of the file usesov_opset. For consistency, the existingov_opsetshould be used. This also applies to usingov_opset.floor_modinstead ofov.mod. - Unused Code: The
last2variable is assigned but never used.
Here is a refactored version of the code that addresses these points for better maintainability.
from .core import _convert_to_node
def diagonal(x, offset=0, axis1=0, axis2=1):
"""OpenVINO backend decomposition for keras.ops.diagonal."""
x_node = _convert_to_node(x) # -> ov.Node
offset_const = ov_opset.constant(int(offset), dtype="i64")
# rank & normalize axes
shape = ov_opset.shape_of(x_node) # i64 vector
rank = ov_opset.shape_of(shape) # scalar i64 (len of shape)
rank_val = ov_opset.squeeze(rank) # [] -> scalar
axis1_node = ov_opset.floor_mod(
ov_opset.add(ov_opset.constant(int(axis1), dtype="i64"), rank_val), rank_val
)
axis2_node = ov_opset.floor_mod(
ov_opset.add(ov_opset.constant(int(axis2), dtype="i64"), rank_val), rank_val
)
# If axis1 == axis2, behavior should match numpy error; Keras tests don't hit this,
# so we skip explicit assert to keep graph-friendly.
# Build permutation to move axis1, axis2 to the end
# perm = [all axes except axis1/axis2 in order] + [axis1, axis2]
arange = ov_opset.range(
ov_opset.constant(0, dtype="i64"), rank_val, ov_opset.constant(1, dtype="i64")
)
mask1 = ov_opset.equal(arange, axis1_node)
mask2 = ov_opset.equal(arange, axis2_node)
not12 = ov_opset.logical_not(ov_opset.logical_or(mask1, mask2))
others = ov_opset.squeeze(
ov_opset.non_zero(not12), [1]
) # gather positions != axis1, axis2
perm = ov_opset.concat(
[others, ov_opset.reshape(axis1_node, [1]), ov_opset.reshape(axis2_node, [1])], 0
)
x_perm = ov_opset.transpose(x_node, perm)
permuted_shape = ov_opset.shape_of(x_perm)
d1 = ov_opset.gather(
permuted_shape,
ov_opset.constant([-2], dtype="i64"),
ov_opset.constant(0, dtype="i64"),
)
d2 = ov_opset.gather(
permuted_shape,
ov_opset.constant([-1], dtype="i64"),
ov_opset.constant(0, dtype="i64"),
)
d1 = ov_opset.squeeze(d1) # scalar
d2 = ov_opset.squeeze(d2) # scalar
# start1 = max(0, offset), start2 = max(0, -offset)
zero = ov_opset.constant(0, dtype="i64")
start1 = ov_opset.maximum(zero, offset_const)
start2 = ov_opset.maximum(zero, ov_opset.negative(offset_const))
# L = min(d1 - start1, d2 - start2)
l1 = ov_opset.subtract(d1, start1)
l2 = ov_opset.subtract(d2, start2)
L = ov_opset.minimum(l1, l2)
# r = range(0, L, 1) -> shape [L]
r = ov_opset.range(zero, L, ov_opset.constant(1, dtype="i64"))
idx_row = ov_opset.add(r, start1)
idx_col = ov_opset.add(r, start2)
idx_row = ov_opset.unsqueeze(
idx_row, ov_opset.constant(1, dtype="i64")
) # [L,1]
idx_col = ov_opset.unsqueeze(
idx_col, ov_opset.constant(1, dtype="i64")
) # [L,1]
diag_idx = ov_opset.concat([idx_row, idx_col], 1) # [L,2]
# Broadcast indices to batch dims: target shape = (*batch, L, 2)
# batch_rank = rank(x) - 2
two = ov_opset.constant(2, dtype="i64")
batch_rank = ov_opset.subtract(rank_val, two)
# build target shape: concat(permuted_shape[:batch_rank], [L, 2])
batch_shape = ov_opset.strided_slice(
permuted_shape,
begin=ov_opset.constant([0], dtype="i64"),
end=ov_opset.reshape(batch_rank, [1]),
strides=ov_opset.constant([1], dtype="i64"),
begin_mask=[0],
end_mask=[0],
)
target_shape = ov_opset.concat(
[batch_shape, ov_opset.reshape(L, [1]), ov_opset.constant([2], dtype="i64")], 0
)
bcast_idx = ov_opset.broadcast(diag_idx, target_shape)
# GatherND with batch_dims = batch_rank
gathered = ov_opset.gather_nd(x_perm, bcast_idx, batch_rank)
return OpenVINOKerasTensor(gathered)
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #21584 +/- ##
==========================================
- Coverage 82.73% 82.68% -0.06%
==========================================
Files 567 567
Lines 56678 56717 +39
Branches 8839 8839
==========================================
+ Hits 46895 46896 +1
- Misses 7609 7647 +38
Partials 2174 2174
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
| d2 = ov_opset.squeeze(d2) # scalar | ||
|
|
||
| # start1 = max(0, offset), start2 = max(0, -offset) | ||
| zero = ov_opset.constant(0, dtype="i64") |
There was a problem hiding this comment.
You can define it once and reuse it throughout the code, since I see the same constant appearing multiple times. This would make the code cleaner and easier to maintain. Please try to clean up the code properly.
| 0, | ||
| ) | ||
| bcast_idx = ov_opset.broadcast(diag_idx, target_shape) | ||
|
|
There was a problem hiding this comment.
If possible, try to avoid broadcast operations, as they tend to increase memory usage.
|
For the format issue: pip install pre-commitThen, run the following command locally to automatically fix formatting issues: pre-commit run --all-files --hook-stage manual |
Summary
Implements support for
numpy.diagonalin the OpenVINO backend for Keras 3.Changes
diagonalop decomposition inkeras/src/backend/openvino/numpy.pyusing OpenVINO opset.diagonalfromexcluded_concrete_tests.txtto enable corresponding tests.Fixes #20910
[Good First Issue][Keras 3 OpenVINO Backend]: Support numpy.diag operation #29115
CC @rkazants