Add new BOUNDS_CHECK GatherScatterMode

[jakevdp]

No description provided.

[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 introduces a new BOUNDS_CHECK mode for JAX's gather and scatter operations, enhancing the robustness and debuggability of array indexing. This mode explicitly raises an IndexError for out-of-bounds static indices, providing immediate feedback to users about invalid access patterns. This prevents silent data manipulation that could occur with other modes like FILL_OR_DROP when indices are outside the array boundaries. The changes involve updating the GatherScatterMode enum, integrating this new behavior into the core indexing logic, and ensuring proper test coverage and compatibility with existing performance models.

Highlights

• New GatherScatterMode: BOUNDS_CHECK: Introduced a new GatherScatterMode called BOUNDS_CHECK. This mode is designed to raise an IndexError when out-of-bounds indices are detected, particularly for static indices or when used with jax.experimental.checkify. For dynamic indices, its behavior defaults to FILL_OR_DROP.
• Integration into Gather and Scatter Operations: The new BOUNDS_CHECK mode has been integrated into the core gather and scatter functions across jax/_src/lax/slicing.py, jax/_src/numpy/indexing.py, and jax/_src/ops/scatter.py to enable its error-raising behavior.
• Static Index Bounds Checking: Logic has been added to jax/_src/numpy/indexing.py to perform explicit static bounds checking. When BOUNDS_CHECK mode is active, index_to_gather will now raise an IndexError if a static index falls outside the array's dimensions.
• Updated Roofline Model: The roofline performance model in jax/experimental/roofline/rooflines.py and its associated tests have been updated to correctly account for the BOUNDS_CHECK mode, treating its FLOPs calculation similarly to FILL_OR_DROP.
• Comprehensive Test Coverage: New parameterized tests have been added in tests/lax_numpy_indexing_test.py to thoroughly validate the BOUNDS_CHECK mode, ensuring that IndexError is correctly raised for various out-of-bounds static indexing scenarios in both gather and scatter operations.

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 new BOUNDS_CHECK mode for GatherScatterMode. This is a valuable addition that will help catch out-of-bounds indexing errors at compile time for static indices. The implementation is well-integrated across the relevant parts of JAX, including lax.slicing, numpy.indexing, and ops.scatter. The fallback behavior to FILL_OR_DROP when static checking is not possible is sensible. The new functionality is also thoroughly tested.

I've made a couple of minor suggestions to use tuples instead of lists for in checks for consistency across the codebase. Overall, this is a great change.