Implement ldexp function in keras.ops #21863
Conversation
Summary of ChangesHello @shashaka, 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 integrates the Highlights
Using Gemini Code AssistThe 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
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 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
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements keras.ops.ldexp for NumPy, TensorFlow, PyTorch, and JAX backends. The overall structure and API definition are good. However, I've found some issues in the backend implementations related to dtype handling and potential precision loss which could lead to incorrect results. I've provided suggestions to improve the correctness and robustness of the implementations for the NumPy, TensorFlow, and PyTorch backends. The tests are well-structured and cover various cases.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #21863 +/- ##
==========================================
- Coverage 82.57% 82.57% -0.01%
==========================================
Files 577 577
Lines 59599 59650 +51
Branches 9351 9356 +5
==========================================
+ Hits 49213 49254 +41
- Misses 7978 7984 +6
- Partials 2408 2412 +4
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:
|
hertschuh
left a comment
hertschuh
left a comment
There was a problem hiding this comment.
Thank you for the addition!
| x1 = convert_to_tensor(x1) | ||
| x2 = convert_to_tensor(x2) | ||
| dtype = dtypes.result_type(x1.dtype, x2.dtype, float) | ||
|
|
||
| if standardize_dtype(x2.dtype) not in dtypes.INT_TYPES: | ||
| raise TypeError( | ||
| f"ldexp exponent must be an integer type. " | ||
| f"Received: x2 dtype={x2.dtype}" | ||
| ) |
There was a problem hiding this comment.
Is this needed on NumPy? Isn't the type promotion already consistent with jax.numpy?
There was a problem hiding this comment.
I may not have fully understood the review at first,
but if the suggestion was to simply return np.ldexp(x1, x2),
I tried that approach and it caused dtype mismatches with JAX in DtypeTest. So I think the explicit dtype handling is still required.
| x1 = tf.cast(x1, tf.float32 if not x1.dtype.is_floating else x1.dtype) | ||
| x2 = tf.cast(x2, x1.dtype) | ||
| result = x1 * tf.pow(tf.constant(2.0, dtype=x1.dtype), x2) | ||
| return tf.cast(tf.where(tf.math.is_inf(x1) | (x1 == 0), x1, result), dtype) |
There was a problem hiding this comment.
Why is the tf.where needed here? What happens otherwise?
There was a problem hiding this comment.
Actually, I referred to the JAX implementation for this and tested the behavior. The where ensures correct results in edge cases. Please see the examples below.
import tensorflow as tf
x1 = tf.constant([0], dtype=tf.float16)
x2 = tf.constant([999], dtype=tf.int32)
result_normal = x1 * (tf.constant(2.0, dtype=tf.float16) ** tf.cast(x2, tf.float16))
result_where = tf.where(tf.math.is_inf(x1) | (x1 == 0), x1, result_normal)
tf.print("Without _where:", result_normal)
tf.print("With _where:", result_where)
Without _where: [-nan]
With _where: [0]
import tensorflow as tf
x1 = tf.constant([np.inf], dtype=tf.float16)
x2 = tf.constant([-99999], dtype=tf.int32)
result_normal = x1 * (tf.constant(2.0, dtype=tf.float16) ** tf.cast(x2, tf.float16))
result_where = tf.where(tf.math.is_inf(x1) | (x1 == 0), x1, result_normal)
tf.print("Without _where:", result_normal)
tf.print("With _where:", result_where)
Without _where: [-nan]
With _where: [inf]
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces the keras.ops.ldexp function, providing implementations for NumPy, TensorFlow, PyTorch, and JAX backends. The changes are well-structured, with corresponding tests for shape inference, correctness, and data types. The implementations for JAX, NumPy, and PyTorch correctly use their respective native ldexp functions. My main feedback is to refactor the TensorFlow backend implementation to use the native tf.math.ldexp function, which would simplify the code and improve its robustness.
google-ml-butler
bot
added
kokoro:force-run
ready to pull
google-ml-butler
bot
removed
awaiting review
ready to pull
5 participants
Adds keras.ops.ldexp, which multiplies a floating-point tensor by 2 raised to the power of an integer tensor, element-wise.
Supported across NumPy, TensorFlow, PyTorch, and JAX backends.
Not supported on OpenVINO.