Weekly GitHub Report for Jax

Thank you for subscribing to our weekly newsletter! Each week, we deliver a comprehensive summary of your GitHub project's latest activity right to your inbox, including an overview of your project's issues, pull requests, contributors, and commit activity.

Table of Contents

• I. News
  • 1.1. Recent Version Releases
  • 1.2. Other Noteworthy Updates
• II. Issues
  • 2.1. Top 5 Active Issues
  • 2.2. Top 5 Stale Issues
  • 2.3. Open Issues
  • 2.4. Closed Issues
  • 2.5. Issue Discussion Insights

• III. Pull Requests
  • 3.1. Open Pull Requests
  • 3.2. Closed Pull Requests
  • 3.3. Pull Request Discussion Insights

• IV. Contributors
  • 4.1. Contributors

I. News

1.1 Recent Version Releases:

The current version of this repository is jax-v0.9.0.1

1.2 Version Information:

On February 3, 2026, JAX v0.9.0.1 was released as a patch update to v0.9.0, incorporating four specific pull requests from the OpenXLA repository to address targeted improvements or fixes without introducing major changes.

Click here to view the full release notes!

II. Issues

2.1 Top 5 Active Issues:

We consider active issues to be issues that that have been commented on most frequently within the last week. Bot comments are omitted.

1. [BUG] Undocumented arguments for jax.vmap: spmd_axis_name and sum_match: This issue addresses the lack of documentation for two arguments, spmd_axis_name and sum_match, in the jax.vmap function. The discussion focuses on whether these arguments should be documented, with spmd_axis_name intentionally left undocumented pending future changes, while sum_match is considered useful enough to warrant documentation.

   • The comments reveal that spmd_axis_name is intentionally undocumented and may be deprecated or reworked, while there is interest in documenting sum_match. A contributor offered to create a documentation PR, which was subsequently submitted to document sum_match with an example, leaving spmd_axis_name undocumented.
   • Number of comments this week: 4

2. [BUG] [Pallas] TPU interpret mode gives wrong result when scalar prefetch SMEM is mutated and pipeline_mode=pl.Buffered(1): This issue describes a problem with the TPU interpret mode in the Pallas library where mutating a scalar prefetch shared memory buffer while using pipeline_mode=pl.Buffered(1) results in incorrect output due to premature computation of block indices. The user expects synchronous single-buffered input handling to delay the next block index computation until after the current kernel invocation, but the interpreter currently performs a lookahead that causes stale data to be used in subsequent kernel executions.

   • The comments acknowledge the bug and clarify that scalar prefetch buffers should be treated as read-only, suggesting documentation updates and potential runtime checks to prevent or warn about modifications; there is also a proposal to introduce a "const ref" mechanism to enforce immutability of these buffers.
   • Number of comments this week: 4

3. [ENHANCEMENT] [P3 (NO SCHEDULE)] jax.lax.betainc does not support gradients with respect to a and b parameters: This issue reports that the function jax.lax.betainc currently supports gradient computation only with respect to the parameter x but raises a ValueError when attempting to differentiate with respect to the shape parameters a or b, limiting its use in training loops where these parameters are learnable. The user requests that gradients with respect to a and b be implemented, noting that closed-form derivatives exist and this is a missing feature rather than a fundamental limitation.

   • The comments clarify that this behavior is intentional due to the lack of easily-expressible closed-form gradients for a and b, marking it as a feature request unlikely to be resolved soon; users confirm the error is explicit in the code, discuss workarounds using jax.scipy.special.betainc which still rely on lax.betainc and thus fail similarly, and highlight the need for proper JVP rules to enable gradient support for these parameters.
   • Number of comments this week: 3

4. [BUG] no_tracing raises on calls that don't trace: the guard at the top of cache_miss fires on any C++-fastpath miss, and functions with host callbacks never take the fastpath: This issue describes a problem where the jax.no_tracing() context manager incorrectly raises an error on calls to jitted functions containing host callbacks, even though no actual tracing occurs, due to a coarse check firing on any C++ dispatch cache miss rather than genuine tracing. The report diagnoses that this behavior contradicts the documented semantics of no_tracing and suggests resolutions including removing the coarse check or renaming it to better reflect its purpose, supported by a minimal reproducible example and detailed analysis.

   • The comments discuss a linked pull request that removes the coarse dispatch-level check causing false positives, explaining that the precise tracing check already covers all genuine cases; the PR adds tests and notes a behavior change where some transformations no longer raise errors, and a side discussion briefly addresses a question about AI-generated comments.
   • Number of comments this week: 3

5. [ENHANCEMENT] [GOOD FIRST ISSUE] [DOCUMENTATION] lax.conv_transpose has different convention than TF, document it and make switching easy?: This issue addresses the difference in kernel spatial axis conventions between JAX's lax.conv_transpose function and TensorFlow's transposed convolution implementation, proposing to document this discrepancy and add an argument to easily switch to the TensorFlow convention. The discussion explores whether such an argument already exists, the usefulness of adding clearer documentation and examples, and the current status of the issue with suggestions for finalizing the documentation.

   • Commenters confirmed that an argument to switch to the TensorFlow convention (transpose_kernel=True) already exists and that the difference is now documented in the function's docstring; several contributors expressed interest in improving documentation and examples, and the issue appears resolved except for a minor docstring wording update.
   • Number of comments this week: 2

2.2 Top 5 Stale Issues:

We consider stale issues to be issues that has had no activity within the last 30 days. The team should work together to get these issues resolved and closed as soon as possible.

As of our latest update, there are no stale issues for the project this week.

2.3 Open Issues

This section lists, groups, and then summarizes issues that were created within the last week in the repository.

Issues Opened This Week: 19

Summarized Issues:

• JAX function output and gradient display issues: The JAX 101 tutorial example for value_and_grad does not display the returned gradients, only printing the loss value twice, which obscures the assigned gradients Wb_grad. Additionally, a regression in JAX versions 0.10.0 and 0.10.1 causes a jitted value_and_grad function for a Flax model on CPU to abort with a static assertion failure, indicating instability in gradient-related functionality.
• [issues/39140, issues/39264]

• GPU and CPU backend compilation and runtime errors: Multiple issues report backend-specific failures including a compile error in the XLA GPU backend due to a tensor size type mismatch in scf.if operations, a regression causing multi-stream windowed einsum operations to fail HLO verification on GPU, and a CPU-specific crash or hang caused by a rank-7/8 broadcast-multiply-reduce operation using XLA's __ynn_fusion. These errors highlight instability and regressions in both GPU and CPU compilation paths.
• [issues/39157, issues/39200, issues/39257]

• SPMD and sharding related compilation failures: Several issues describe SPMD partitioner failures such as malformed concatenate operations during reverse-mode differentiation through jax.linear_transpose on fully sharded axes, and a regression where reshaping arrays with size-1 explicit mesh axes causes axis dropping and sharding mismatches. These problems cause HLO verifier errors and sharding inconsistencies in distributed computations.
• [issues/39301, issues/39309]

• Pallas library and shared memory buffer issues: The Pallas TPU interpret mode exhibits incorrect results when mutating a scalar prefetch shared memory buffer combined with buffered pipeline mode, due to premature computation of block indices. Additionally, the pl.empty_ref_like function does not accept existing Ref objects, complicating usage by requiring redundant shape and dtype specifications.
• [issues/39179, issues/39190]

• vmap transformation and reference (state) effects problems: Using vmap over Ref arguments with non-zero in_axes causes shape mismatches and attribute errors, and support for reference effects within vmap of lax.cond is currently missing, raising NotImplementedError. These issues limit the composability of stateful operations with vectorized mapping.
• [issues/39159, issues/39288, issues/39290]

• Tracing and dispatch cache errors: The jax.no_tracing() context manager incorrectly raises errors on jitted functions with host callbacks due to a coarse cache miss check in the C++ dispatch fallback, leading to misleading error messages and unnecessary overhead despite no actual tracing occurring.
• [issues/39289]

• FFT and linear algebra failures on specific hardware: A compilation failure occurs in XLA when performing a jitted complex128 FFT with a sharded transform axis due to mixed-type multiplication errors from complex64 twiddle factors. Also, batched jnp.linalg.eigh crashes on many-core CPUs because of OpenBLAS thread oversubscription, causing heap corruption with large batches of small symmetric matrices.
• [issues/39291, issues/39292]

• Key flattening and data structure handling bugs: The flatten_one_level_with_keys function incorrectly assigns the first namedtuple field key to all child elements, causing inconsistent behavior compared to full keyed flattening and potentially leading to errors in keyed data processing.
• [issues/39297]

• Performance improvements for random number generation: There are proposals to improve the performance of jax.random.poisson by replacing its loop-based implementation with a faster approximate arithmetic-only sampler, and to enhance other loop-based random number generators for nontrivial distributions by implementing loop-free or faster sampling methods, aiming to increase throughput and reduce latency on GPUs.
• [issues/39312, issues/39314]

2.4 Closed Issues

This section lists, groups, and then summarizes issues that were closed within the last week in the repository. This section also links the associated pull requests if applicable.

Issues Closed This Week: 6

Summarized Issues:

• Performance improvements in sampling methods: This topic covers the implementation of a faster, approximate gamma distribution sampler for JAX based on Goldstein's series expansion algorithm, designed as an alternative to the slower exact rejection sampler. The goal is to improve performance while allowing users to opt-in to the new method.
• issues/38141

• Type errors and compatibility issues with Python versions: This topic involves a TypeError encountered when building jaxlib 0.10.2 with Python 3.10, caused by an unexpected 'required' keyword argument in the argparse module's _SubParsersAction.init() method. This indicates compatibility problems between JAX and newer Python versions.
• issues/38969

• Numerical precision and data type conversion bugs: This topic addresses a bug where converting the smallest subnormal value of jnp.bfloat16 to jnp.float64 and back results in zero instead of preserving the original value. This highlights issues in maintaining precision during type conversions in JAX.
• issues/38995

• Incorrect handling of NaN values in numpy functions: This topic reports that jax.numpy.heaviside incorrectly returns the second input value when the first input is NaN, instead of propagating NaN as NumPy does. This discrepancy indicates inconsistent behavior with standard NumPy in handling NaN inputs.
• issues/39109

• Memory management and profiling regressions: This topic covers a change in the behavior of the XLA_PYTHON_CLIENT_ALLOCATOR=platform flag in recent JAX versions, where memory is no longer deallocated as expected during profiling. This leads to increased and fragmented GPU memory usage, differing from earlier versions' behavior.
• issues/39145

• Multi-GPU communication errors due to conflicting libraries: This topic describes failures in multi-GPU NCCL collectives with a "corrupted comm object detected" error when TensorFlow is imported alongside JAX in the same process. The issue arises from multiple NCCL implementations being loaded simultaneously, causing communicator incompatibility on Kaggle T4 x2 setups.
• issues/39167

2.5 Issue Discussion Insights

This section will analyze the tone and sentiment of discussions within this project's open and closed issues that occurred within the past week. It aims to identify potentially heated exchanges and to maintain a constructive project environment.

Based on our analysis, there are no instances of toxic discussions in the project's open or closed issues from the past week.

III. Pull Requests

3.1 Open Pull Requests

This section provides a summary of pull requests that were opened in the repository over the past week. The top three pull requests with the highest number of commits are highlighted as 'key' pull requests. Other pull requests are grouped based on similar characteristics for easier analysis. Up to 25 pull requests are displayed in this section, while any remaining pull requests beyond this limit are omitted for brevity.

Pull Requests Opened This Week: 14

Key Open Pull Requests

1. [oneAPI] Add CI workflow to build hermetic JAX wheel artifacts: This pull request introduces a dedicated oneAPI continuous integration workflow to build hermetic JAX wheel artifacts as part of the JAX presubmit process.

• URL: pull/39282

• Associated Commits: c8ca6, 3367e, e7d64

2. Improve pytree prefix mismatch errors involving None specs and dummy …: This pull request improves error handling for pytree prefix mismatches in JAX by correctly treating None as a leaf in broadcasting checks, preventing false mismatch reports and internal type leaks, and clarifying error messages to better distinguish between leaves and containers with no leaves.

• URL: pull/39239

• Associated Commits: a9064, 075d1

3. Raise a clear error for unsupported dtypes in LAPACK-backed linalg on CPU: This pull request improves error handling in LAPACK-backed linear algebra operations on CPU by adding explicit dtype checks that raise clear, actionable NotImplementedError messages when unsupported float16 or bfloat16 inputs are used, replacing previously confusing internal errors.

• URL: pull/39295

• Associated Commits: 00e91, 6f92e

Other Open Pull Requests

• Documentation fixes for jax.value_and_grad examples: This pull request corrects the examples for jax.value_and_grad by updating them to print the computed gradients instead of printing the loss value twice. This ensures the documentation accurately demonstrates how gradients are returned alongside function values.
  • pull/39160

• Hermetic dependencies and OneAPI support: These pull requests add runfiles data files and ROCm hermetic dependencies to the plugin and PJRT wheels for proper hermetic Bazel test execution, and also add OneAPI continuous integration test scripts and infrastructure to support Intel GPU testing. Additionally, support for OneAPI buffer donation is introduced by appending oneapi to the list of platforms allowing buffer donation.
  • pull/39220, pull/39285, pull/39303

• Fixes and clarifications for lax.conv_transpose transpose_kernel docstring: These pull requests clarify and neutrally reword the docstring for the transpose_kernel parameter in lax.conv_transpose to resolve contradictions with the Notes section about TensorFlow/Keras compatibility. They also remove a contradictory sentence labeling the parameter as "completely pointless," improving documentation clarity and fixing markup issues.
  • pull/39227, pull/39230

• Bug fix for jax.scipy.special.betaln behavior with infinite inputs: This pull request fixes the jax.scipy.special.betaln function to correctly return -inf instead of nan when one argument is infinite and the other is a finite positive number. It aligns the function's behavior with SciPy by adding a guard after argument swapping to properly handle infinite inputs.
  • pull/39224

• Improved error handling for out-of-range 64-bit integer inputs: This pull request modifies JAX to raise a Python OverflowError when 64-bit NumPy integer inputs contain out-of-range values while jax_enable_x64 is disabled. This replaces the previous unsafe silent truncation to 32 bits, ensuring consistent error handling and improving safety in integer downcasting.
  • pull/39273

• Fix for prime factorization bug in mesh_utils: This pull request fixes a bug in the _get_prime_factors function where a for…else loop without a break caused smaller prime factors to be dropped if the largest prime factor exceeded isqrt(x) + 1. The fix includes a two-line code change and extended tests to cover previously failing cases, restoring correct logical mesh construction with allow_split_physical_axes=True.
  • pull/39300

• Fix for false-positive jax.no_tracing() check in JIT dispatch: This pull request removes an incorrect check for jax.no_tracing() in the JIT dispatch slow path that raised errors on warm calls of jitted functions with host callbacks. The overly broad check in cache_miss is deleted, ensuring only genuine re-traces trigger errors, and a test is added to verify this behavior.
  • pull/39302

• Fix for namedtuple keys in one-level pytree flattening: This pull request addresses issue #39297 by fixing the handling of namedtuple keys in the one-level pytree flattening process. The fix ensures correct functionality and passes relevant tests.
  • pull/39306

3.2 Closed Pull Requests

This section provides a summary of pull requests that were closed in the repository over the past week. The top three pull requests with the highest number of commits are highlighted as 'key' pull requests. Other pull requests are grouped based on similar characteristics for easier analysis. Up to 25 pull requests are displayed in this section, while any remaining pull requests beyond this limit are omitted for brevity.

Pull Requests Closed This Week: 40

Key Closed Pull Requests

1. JAX-PRESUBMIT TEST: This pull request is about adding and updating presubmit test configurations for JAX, specifically involving changes to the bazel_oneapi_presubmit.yml file and other related test updates.

• URL: pull/39169

• Associated Commits: 806e8, 95150, 2049c

• Associated Commits: 806e8, 95150, 2049c

2. [docs] add 501 (systems topics); add CuTe DSL to 401 and errors to 101: This pull request updates the documentation by adding a new 501 section covering advanced systems topics such as multi-controller JAX, distributed data loading, fault tolerance, exporting and serialization, and related utilities, while also incorporating the CuTe DSL tutorial into the 401 section, adding an errors reference to 101, updating cross-links from 201 to 501, and enhancing the new_docs landing page.

• URL: pull/39095

• Associated Commits: 79502, 38083

• Associated Commits: 79502, 38083

3. [hijax] fix linearize bug in presence of non-differentiable args: This pull request fixes a bug in the hijax linearize function that occurred when non-differentiable arguments were present.

• URL: pull/39127

• Associated Commits: 5d31e, a6ca7

• Associated Commits: 5d31e, a6ca7

Other Closed Pull Requests

• JAX 0.11.0 Release and Documentation Updates: This topic covers the finalization and official release of JAX version 0.11.0 along with updates to the changelog and supported versions of NumPy and SciPy. The pull requests ensure accurate documentation of deprecations and compatibility for the new release.
  • pull/39241, pull/39151, pull/39148

• GPU Device Visibility and Environment Variable Management: These pull requests improve GPU device visibility handling in pytest xdist workers by rotating device visibility and changing environment variable assignment methods. The changes prevent unintended GPU sharing and distribute test loads evenly across GPUs.
  • pull/36806, pull/36840

• HIP Memory Leak Fixes in RNN Kernels: This set of changes fixes HIP memory leaks by properly freeing GPU memory allocated in RNN kernel paths and adding missing free macros. These fixes address resource management issues in the HIP backend.
  • pull/35638

• OneAPI GPU Plugin Enhancements: These pull requests add dlpack array interoperability support for the OneAPI GPU plugin and enable JAX oneAPI wheels to bundle and dynamically load their own oneAPI 2025.1 runtime libraries. This removes the dependency on system-installed oneAPI components and improves compatibility.
  • pull/38765, pull/38815

• Random Gamma Sampling Improvements: A new method parameter is introduced to jax.random.gamma and jax.random.loggamma functions, allowing users to select between an exact rejection sampler and a faster, fully vectorized approximate sampler. The approximate sampler offers differentiable gradients and better accelerator performance at the cost of a small bias.
  • pull/38358

• Array API Version Update and Compatibility Fixes: The array API version is updated to 2025.12, which introduces minor breaking changes such as returning tuples instead of lists. Temporary test skips are added to handle incompatibilities, with plans for individual fixes in future pull requests.
  • pull/39073, pull/39133, pull/39155

• VJP Object Enhancements and Documentation: These pull requests introduce a saveable_args keyword argument to jax.vjp for selective saving of arguments, improving serialization and offloading of VJP objects. Additionally, detailed documentation on first-class VJP objects is added, reorganizing related content and updating indexes accordingly.
  • pull/39104, pull/39206, pull/39201

• XLA Metadata API Improvements: The implementation of xla_metadata_call is moved for better organization, and a new function xla_metadata_call2 is introduced with enhanced metadata handling. The related documentation is rewritten to improve API robustness and ensure correct metadata propagation through autodiff computations.
  • pull/39156, pull/39201

• ROCm GPU Compatibility Fix: A fallback mechanism is implemented to resolve GPU architecture directly from the device's ISA when the device marketing name cannot be matched to a GpuVersion. This fix prevents failures in get_gpu_info() and enables Triton calls to recognize supported GPU devices on ROCm.
  • pull/39038

• Error Message and Test Improvements: Improvements include replacing an opaque TypeError with a clear NotImplementedError for missing GSPMD sharding propagation rules, adding a test for output forwarding in JIT rematerialization, fixing a corner case in the checkify module, simplifying TPU test debug print calls, and adding a flaky test to the CI ignore list.
  • pull/39128, pull/39132, pull/39147, pull/39212

• Build Script Compatibility Fix: The build script is modified to set the required flag for add_subparsers as an attribute instead of a keyword argument. This change ensures compatibility with third-party argparse backports that do not support the required keyword introduced in Python 3.7.
  • pull/39001

3.3 Pull Request Discussion Insights

This section will analyze the tone and sentiment of discussions within this project's open and closed pull requests that occurred within the past week. It aims to identify potentially heated exchanges and to maintain a constructive project environment.

Based on our analysis, there are no instances of toxic discussions in the project's open or closed pull requests from the past week.

IV. Contributors

4.1 Contributors

Active Contributors:

We consider an active contributor in this project to be any contributor who has made at least 1 commit, opened at least 1 issue, created at least 1 pull request, or made more than 2 comments in the last month.

If there are more than 10 active contributors, the list is truncated to the top 10 based on contribution metrics for better clarity.