Fix Ref indexing discrepancies with array indexing

[yashwantbezawada]

This addresses the issues raised in #33322 where Ref indexing behaves differently from regular JAX array indexing.

Changes:

• Fixed OOB slice validation to allow empty slices (e.g., slice(11, 12, 1) on size-10 array now returns empty result instead of erroring)
• Added support for None indexing to insert new axes (e.g., x[None], x[..., None])
• Added support for negative step slices (e.g., x[::-1]) by converting them to positive steps and applying a flip transform
• Refactored RefIndexer.__getitem__ into helper functions for better maintainability

For None indexing, added a RefNewAxis transform that tracks where new axes should be inserted. For negative steps, added a RefFlip transform that reverses the specified axes after indexing with the positive-step equivalent.

Fixes #33322.

[gemini-code-assist]

Summary of Changes

Hello @yashwantbezawada, 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 aims to improve the consistency and robustness of Ref indexing by aligning its behavior more closely with standard JAX array indexing. It resolves issues related to slice validation, introduces support for inserting new dimensions, and provides clearer feedback for unsupported operations, making Ref indexing more intuitive and less prone to unexpected errors.

Highlights

• Out-of-Bounds Slice Validation: Refined the validation logic for out-of-bounds slices to correctly handle and return empty slices (e.g., slice(11, 12, 1) on a size-10 array) instead of raising an error.
• Support for None Indexing: Added functionality to support None indexing (equivalent to np.newaxis) which allows users to insert new dimensions (axes) into the array's shape, such as x[None] or x[..., None].
• Improved Error Message for Negative Slice Steps: Enhanced the error message for negative slice steps to clearly state that Ref indexing does not support them, providing a suggestion to use positive steps and reverse the result if needed.

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.

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

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 several fixes to make Ref indexing more consistent with JAX array indexing. It correctly handles out-of-bounds empty slices, adds support for None to insert new axes, and improves the error message for negative slice steps. The changes are well-implemented. I have one suggestion to refactor a piece of duplicated logic to improve code maintainability.

[yashwantbezawada]

Updated this PR to fully fix all the issues from #33322.

Added RefFlip transform to handle negative slice steps - these get converted to positive equivalents at the RefIndexer level and then flipped to get the right element order. Also updated discharge.py to handle both RefNewAxis and RefFlip in the transform functions.

Refactored the duplicated code in RefIndexer.getitem as suggested in the review.

All four test cases from the issue should work now (OOB clamping, None indexing, ellipsis+None, and negative slices).

[gemini-code-assist]

Code Review

This pull request is a great step towards aligning Ref indexing with standard JAX array indexing. The changes to allow empty slices, support None for new axis insertion, and improve the error message for negative steps are all valuable improvements. The implementation using RefNewAxis and RefFlip transforms is a clean and effective approach. My review includes a few suggestions to refactor some loops into single function calls for better performance and to break down a particularly complex method for improved readability. Overall, this is a solid contribution.