Update requirements lock files for JAX 0.8.1.

[danielsuo]

No description provided.

[gemini-code-assist]

Summary of Changes

Hello @danielsuo, 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 focuses on updating the project's dependency lock files to align with JAX 0.8.1. The changes primarily involve significant revisions to JAX CUDA dependencies, including a version bump for CUDA 13 plugins and the removal of CUDA 12 support. Additionally, the libtpu dependency has been dropped, and a large set of explicit NVIDIA CUDA library pins have been removed, streamlining the overall dependency management for GPU acceleration.

Highlights

• JAX CUDA Dependency Update: The JAX CUDA-related packages have been updated across all lock files. Specifically, jax-cuda12-pjrt and jax-cuda12-plugin at version 0.8.0 have been removed, and jax-cuda13-pjrt and jax-cuda13-plugin have been updated from 0.8.0 to 0.0.1rc0. This indicates a shift towards CUDA 13 and potentially a pre-release version of its JAX plugins.
• Removal of libtpu: The libtpu==0.0.24 dependency has been entirely removed from all requirements lock files, suggesting it is no longer a required component for the project.
• Streamlined NVIDIA CUDA Dependencies: A significant number of explicit NVIDIA CUDA library dependencies, including various nvidia-cublas, nvidia-cuda-crt, nvidia-cudnn, nvidia-cufft, nvidia-cusolver, nvidia-cusparse, nvidia-nccl, nvidia-nvjitlink, and nvidia-nvvm packages, have been removed from the lock files. This change likely simplifies CUDA dependency management, possibly by relying on JAX's internal handling or system-wide CUDA installations.
• Minor Dependency Comment Updates: Small adjustments were made to comments within the lock files, such as removing the # wheel comment from the packaging dependency and clarifying the source for pytest to include build/test-requirements.txt.

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 updates the requirements lock files. However, there are several major and potentially breaking changes that are not explained. All nvidia-* packages have been removed, which suggests a fundamental change in how CUDA dependencies are handled. This is a critical change that will likely break setups for users relying on JAX with NVIDIA GPUs. libtpu, jax-cuda12-plugin, and jax-cuda12-pjrt dependencies have been dropped from the lock files, despite being present in build/requirements.in. This indicates a dependency resolution failure or an undocumented decision to drop support for TPUs and CUDA 12. jax-cuda13-plugin and jax-cuda13-pjrt have been downgraded from version 0.8.0 to 0.0.1rc0, a release candidate. This is unusual and risky. These changes need to be clarified and justified. If they are intentional, the reasoning should be documented in the pull request description.

[danielsuo]

/gemini review

[gemini-code-assist]

Code Review

This pull request updates the dependency lock files for Python versions 3.11, 3.12, 3.13, and 3.14, including their free-threaded variants. The main purpose is to update JAX and its associated packages to version 0.8.1.

The changes are consistent across all modified lock files. Most dependencies are updated to new patch or minor versions, which is expected during routine maintenance. One notable change is the major version upgrade of wrapt from 1.17.3 to 2.0.1, pulled in as a dependency of tensorflow. This major version of wrapt mainly removes support for legacy Python versions and is not expected to cause issues.

The updates seem correct and well-generated. I have no further comments.