Skip to content

Remove chex dependency#5295

Merged
copybara-service[bot] merged 2 commits into
google:mainfrom
samanklesaria:remove_chex
Mar 4, 2026
Merged

Remove chex dependency#5295
copybara-service[bot] merged 2 commits into
google:mainfrom
samanklesaria:remove_chex

Conversation

@samanklesaria

@samanklesaria samanklesaria commented Mar 2, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

A deprecation warning in chex led to us pinning a numpy version (see #5162). But chex was barely used anyways: this PR replaces the uses of chex and removes the dependency. This fixes #5288

@gemini-code-assist

gemini-code-assist Bot commented Mar 2, 2026

Copy link
Copy Markdown
Contributor

Summary of Changes

Hello, 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 addresses a deprecation warning originating from the chex library, which previously necessitated pinning the numpy version. Given that chex was minimally utilized, this change removes the dependency entirely. The core impact is the restoration of flexible numpy versioning and the internal refactoring of test utilities to remove chex-specific calls, thereby streamlining the project's dependencies and improving maintainability.

Highlights

  • Dependency Removal: The chex dependency has been completely removed from the project, including its presence in pyproject.toml.
  • Numpy Version Unpinning: The temporary numpy version pinning, which was a workaround for a chex deprecation warning, has been removed, allowing for broader numpy compatibility.
  • Test Utility Replacements: All usages of chex-specific test utilities, such as chex.TestCase, chex.assert_type, chex.assert_max_traces, and chex.assert_trees_all_close, have been replaced with equivalent standard unittest assertions or custom implementations.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog
  • pyproject.toml
    • Removed numpy version pinning and associated comments.
    • Removed chex from the testing optional dependencies.
  • tests/jax_utils_test.py
    • Removed the chex import.
    • Changed PadShardUnpadTest to inherit from parameterized.TestCase instead of chex.TestCase.
    • Removed the tearDown method that cleared chex trace counters.
    • Introduced a custom assert_max_traces decorator to replace chex.assert_max_traces.
    • Replaced chex.assert_type calls with self.assertEqual for type assertions.
  • tests/linen/linen_test.py
    • Removed the chex import.
    • Replaced chex.assert_trees_all_close with a jax.tree.map combined with np.testing.assert_allclose for tree comparisons.
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[bot]
gemini-code-assist Bot reviewed Mar 2, 2026
View reviewed changes

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request successfully removes the chex dependency, addressing a numpy versioning conflict. The changes are well-executed, replacing chex functionalities with standard library equivalents or custom helpers across pyproject.toml and test files. I have one suggestion to enhance the new assert_max_traces decorator for better maintainability and debugging.

Comment thread tests/jax_utils_test.py
@samanklesaria @gemini-code-assist
Update tests/jax_utils_test.py
afc169e 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
@samanklesaria samanklesaria requested a review from IvyZX March 2, 2026 20:19
@samanklesaria samanklesaria marked this pull request as ready for review March 2, 2026 20:19
@copybara-service copybara-service Bot merged commit bbfab0a into google:main Mar 4, 2026
21 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@IvyZX IvyZX IvyZX approved these changes

+1 more reviewer

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

pull ready

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

FR: Unpin the numpy dependency upper bound version

2 participants

@samanklesaria @IvyZX