Refactor indexing code

[jakevdp]

The goal here is to reduce duplicated logic at trace time (e.g. attempt_rewriting_take_via_slice andindex_to_gather), to make the indexing logic more extensible (e.g. the recently-added internal-for-now strategy argument), and to make JAX's indexing logic more easily reusable for e.g. pallas indexing.

The approach:

• parse indices only once, and store the results in an NDIndexer structure. This is particularly important for handling ellipses, which requires multiple passes through the raw indices to determine the number of dimensions they consume.
• give this structure methods for commonly-needed operations, such as replacing ellipses, processing boolean masks, handling Nones, validating bounds, etc.
• change existing indexing methods to pass this structure rather than re-parsing raw indices at each phase.

[gemini-code-assist]

Summary of Changes

Hello @jakevdp, 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 undertakes a significant refactoring of JAX's array indexing system. The primary goal is to establish a more robust and maintainable foundation for handling diverse NumPy-style indexing patterns by centralizing the logic for parsing and validating indices. This change aims to improve the clarity of indexing operations, provide more precise error feedback for invalid inputs, and streamline future enhancements to JAX's array manipulation capabilities.

Highlights

• New Indexing Utilities: Introduced IndexType enum, ParsedIndex NamedTuple, and parse_indices function to standardize and centralize the parsing and validation of NumPy-style array indices.
• Refactored Index Validation: The validate_static_indices function has been refactored to utilize the new parse_indices utility, leading to more consistent and robust index boundary checks.
• Improved Error Handling: Enhanced error messages and validation for various invalid indexing scenarios, including out-of-bounds integer indices and incorrect index types.

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 a significant and beneficial refactoring of the NumPy indexing logic. By creating a structured parsing mechanism with IndexType, ParsedIndex, and parse_indices, the validate_static_indices function is greatly simplified and made more modular and readable. The new approach is well-designed and improves the maintainability of the indexing code. I have one minor suggestion to improve clarity in the new bounds checking logic.

[jakevdp]

/gemini review

[gemini-code-assist]

Code Review

This is an excellent and significant refactoring of the indexing logic. The introduction of the NDIndexer class is a great design choice that centralizes complex index parsing and manipulation, making the code cleaner, more extensible, and easier to reason about. The way you've updated rewriting_take and _gather to leverage this new class is very well done. I've found one potential bug in compute_via_static_slice that could lead to an UnboundLocalError, but otherwise, the changes look solid.