Fix unravel_index coordinate order in TensorFlow backend

[buildwithsuhana]

Description

This PR resolves a discrepancy in ops.unravel_index between the TensorFlow backend and NumPy (and other Keras backends). Previously, the TF backend implementation had separate logic for 1D and non-1D indices. The non-1D logic incorrectly iterated through the shape in forward order while calculating coordinates, leading to inverted results for asymmetric shapes.

Example of the bug:

   shape = (4, 5, 6)
   index = 100
   # Expected (NumPy): [3, 1, 4]
   # Actual (TF 0D/2D): [5, 0, 0]  # Incorrectly calculated as if shape was reversed

reproduction script: https://colab.research.google.com/drive/1SdZRRVDUV3Dg8ORpTQmlOIA_U4f3_yyo?resourcekey=0-7XMyFCvdBdNJLhrdHE9Whw&usp=sharing

Contributor Agreement

Please review our AI-Assisted Contribution Policy and check all boxes below before submitting your PR for review:

• I am a human, and not a bot.
• I will be responsible for responding to review comments in a timely manner.
• I will work with the maintainers to push this PR forward until submission.

Note: Failing to adhere to this agreement may result in your future PRs no longer being reviewed.

[gemini-code-assist]

Code Review

This pull request simplifies and corrects the unravel_index implementation in the TensorFlow backend. It removes the redundant conditional block for 1D indices and the incorrect forward-loop logic for multi-dimensional indices, replacing them with a single, unified loop that iterates over the shape in reverse order. This correctly handles multi-dimensional inputs and asymmetric shapes without requiring explicit reshaping. Additionally, corresponding unit tests for asymmetric shapes and multi-dimensional inputs have been added to verify correctness against NumPy. I have no feedback to provide as the changes are clean, correct, and well-tested.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 85.24%. Comparing base (7518459) to head (c0f5a51).
⚠️ Report is 13 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #23434      +/-   ##
==========================================
+ Coverage   84.90%   85.24%   +0.33%     
==========================================
  Files         468      468              
  Lines       70651    73037    +2386     
  Branches    11708    12248     +540     
==========================================
+ Hits        59986    62258    +2272     
- Misses       7658     7722      +64     
- Partials     3007     3057      +50     

Flag	Coverage Δ
keras	85.05% <100.00%> (+0.34%)	⬆️
keras-cpu	83.97% <100.00%> (-0.01%)	⬇️
keras-gpu	69.76% <100.00%> (-0.03%)	⬇️
keras-jax	59.56% <0.00%> (+1.14%)	⬆️
keras-numpy	53.89% <0.00%> (+<0.01%)	⬆️
keras-openvino	59.61% <0.00%> (-0.04%)	⬇️
keras-tensorflow	61.12% <100.00%> (+1.08%)	⬆️
keras-torch	60.63% <0.00%> (+1.11%)	⬆️
keras-tpu	57.22% <0.00%> (-0.04%)	⬇️

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

☔ View full report in Codecov by Harness.
📢 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.

[divyashreepathihalli]

@jules can you verify the changes in this PR?

[github-actions]

Jules

I have verified the changes in this PR.

The PR successfully fixes a bug in unravel_index for non-1D inputs in the TensorFlow backend and correctly matches NumPy's behavior.

Analysis:

1. Bug Fix (Coordinate Order): The previous implementation had a special branch for indices.ndim == 1 that correctly iterated over reversed(shape). However, for inputs of any other dimension (scalars, 2D arrays, etc.), it incorrectly iterated over shape in the forward direction. This caused the dimensions to be calculated in the wrong order, resulting in incorrect coordinate values compared to numpy.unravel_index. The PR fixes this by uniformly iterating over reversed(shape) for all input shapes.
2. Code Simplification: The PR correctly removes the tf.reshape(..., indices_shape) call. Because indices is a TensorFlow tensor, arithmetic operations like % dim and // dim natively preserve the tensor's shape. The explicit reshape was unnecessary, and removing it allows for the removal of the 1D vs N-D conditional branches, greatly simplifying the function.
3. Tests: The PR adds appropriate tests in keras/src/ops/numpy_test.py that check both a scalar input (x = np.array(100)) and a 2D array input (x = np.array([[100, 101]])). These tests effectively cover the previously buggy behavior and guarantee it matches NumPy.

The fix is mathematically correct, simplifies the codebase, and is well-tested.

---

Asked by @divyashreepathihalli · session sessions/8010471067371398690