Add unique op

[pass-lin]

Provide an implementation of keras.ops.numpy.unique
Referencing:
https://docs.jax.dev/en/latest/_autosummary/jax.numpy.unique.html
https://pytorch.org/docs/stable/generated/torch.unique.html

Include backends for TensorFlow, JAX, PyTorch, and NumPy.
However, do not include the OpenVINO backend for now.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! 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 significantly enhances Keras's numerical operations by integrating a unique function into keras.ops.numpy. This addition allows users to efficiently extract unique elements from tensors, a fundamental operation in data processing and machine learning. The implementation ensures broad compatibility across Keras's core backends (TensorFlow, JAX, PyTorch, NumPy), providing a unified and consistent API for this common task. This change improves the framework's utility for a wider range of data manipulation and analysis workflows.

Highlights

• New unique operation: Introduced keras.ops.numpy.unique to find unique elements in tensors, mirroring functionality in JAX and PyTorch.
• Multi-backend support: Provided implementations for TensorFlow, JAX, PyTorch, and NumPy backends, ensuring consistent behavior across Keras's supported execution environments.
• OpenVINO exclusion: Explicitly marked the unique operation as not supported for the OpenVINO backend at this time, raising a NotImplementedError.
• Comprehensive testing: Added a dedicated test suite (UniqueTest) covering various scenarios, including 1D and multi-dimensional arrays, return_inverse, return_counts, axis handling, symbolic tensors, and edge cases like empty arrays or all-same elements.
• API and internal implementation: Included both a functional keras.ops.unique API and an internal Unique Operation class for symbolic execution, complete with detailed docstrings and examples.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

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 by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

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 pull request 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.

Footnotes

1. 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. ↩

[gemini-code-assist]

Code Review

This pull request introduces a new unique operation to Keras's ops.numpy module, providing implementations for various backends (JAX, NumPy, TensorFlow, Torch) and marking it as not implemented for OpenVINO. The change also includes a new Unique operation class, comprehensive docstrings, and extensive unit tests. Feedback indicates a critical bug in the TensorFlow backend's lexicographical sort, which uses a Python for loop that will fail in graph mode; tf.linalg.lexsort is suggested as a fix. Additionally, compatibility issues with the sorted argument in older JAX and NumPy versions are noted, recommending its removal for broader support. Lastly, the input argument in keras.ops.unique should be renamed to x for API consistency.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 97.59036% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 83.30%. Comparing base (c042656) to head (41097d0).
⚠️ Report is 55 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/api/_tf_keras/keras/ops/__init__.py	0.00%	1 Missing ⚠️
keras/api/_tf_keras/keras/ops/numpy/__init__.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #22546      +/-   ##
==========================================
+ Coverage   83.26%   83.30%   +0.04%     
==========================================
  Files         596      596              
  Lines       67815    68172     +357     
  Branches    10558    10621      +63     
==========================================
+ Hits        56463    56792     +329     
- Misses       8608     8636      +28     
  Partials     2744     2744              

Flag	Coverage Δ
keras	83.11% <97.59%> (+0.03%)	⬆️
keras-jax	59.65% <49.39%> (-0.18%)	⬇️
keras-numpy	55.34% <50.60%> (+1.24%)	⬆️
keras-openvino	53.29% <16.86%> (+1.86%)	⬆️
keras-tensorflow	61.07% <86.74%> (-0.06%)	⬇️
keras-torch	59.84% <54.21%> (-0.17%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[hertschuh]

@pass-lin

Thanks for working on this! I had actually started working on this a long time ago and didn't finish. I think you can reuse a lot of it though:
master...hertschuh:keras:unique_op

Some key comments:

• Please add the return_index argument.
• Keep the sorted argument (mine didn't have it)
• I'd like to use the JAX signature, i.e. add the size and fill_value arguments. This is to make it compilable with JAX. Making it compilable makes it way more useful in general. The downside is that it adds complexity for the other backend as you need to pad the result (I have a working implementation for the NumPy backend).

Also please use the current pattern for unit tests instead of having a separate class.

[pass-lin]

@pass-lin

Thanks for working on this! I had actually started working on this a long time ago and didn't finish. I think you can reuse a lot of it though: master...hertschuh:keras:unique_op

Some key comments:

• Please add the return_index argument.
• Keep the sorted argument (mine didn't have it)
• I'd like to use the JAX signature, i.e. add the size and fill_value arguments. This is to make it compilable with JAX. Making it compilable makes it way more useful in general. The downside is that it adds complexity for the other backend as you need to pad the result (I have a working implementation for the NumPy backend).

Also please use the current pattern for unit tests instead of having a separate class.

If we only consider the parameters of the interface,
we can observe that torch.unique ⊆ numpy.unique (version > 2.3) ⊆ jnp.unique.
Therefore, using torch.unique as the standard minimizes bugs, as we only need to adjust the parameter inputs and outputs appropriately.
In contrast, tf.unique is relatively primitive, requiring most functionality to be implemented manually. Consequently, a more complex interface is also more prone to bugs.

Therefore, I recommend using torch.unique as the standard API.

[hertschuh]

@pass-lin
Thanks for working on this! I had actually started working on this a long time ago and didn't finish. I think you can reuse a lot of it though: master...hertschuh:keras:unique_op
Some key comments:

• Please add the return_index argument.
• Keep the sorted argument (mine didn't have it)
• I'd like to use the JAX signature, i.e. add the size and fill_value arguments. This is to make it compilable with JAX. Making it compilable makes it way more useful in general. The downside is that it adds complexity for the other backend as you need to pad the result (I have a working implementation for the NumPy backend).

Also please use the current pattern for unit tests instead of having a separate class.

If we only consider the parameters of the interface, we can observe that torch.unique ⊆ numpy.unique (version > 2.3) ⊆ jnp.unique. Therefore, using torch.unique as the standard minimizes bugs, as we only need to adjust the parameter inputs and outputs appropriately. In contrast, tf.unique is relatively primitive, requiring most functionality to be implemented manually. Consequently, a more complex interface is also more prone to bugs.

Therefore, I recommend using torch.unique as the standard API.

Ok, let's say we remove the return_index argument.

I still want the size and fill_value arguments, otherwise, it makes unique unusable on JAX. The padding logic that need to be added would be very similar for Torch and Numpy and is already done in my branch.