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:

Released on February 3, 2026, JAX v0.9.0.1 is essentially the same as v0.9.0 but includes patches from four specific pull requests, indicating targeted updates and bug fixes rather than major new features.

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] jit silently miscompiles alpha * (A @ B + bias) on sm_120: alpha is applied to the matmul product but not to the bias: This issue reports a silent miscompilation in JAX's GPU JIT compilation on NVIDIA RTX 5090 (sm_120) where a scalar multiplier alpha is applied to the matrix multiplication product but not to the bias term in the fused GEMM epilogue, causing a significant and deterministic error in the output. The problem arises specifically at the default autotune level 4 or higher, leading to large relative errors that are not present in eager mode or CPU backend, and the root cause is traced to a missing epilogue check in XLA's GemmRewriter::HandleMultiply function, which incorrectly folds the scalar multiplier only onto the matrix product but not the bias addition.

   • Comments confirm reproducibility across different GPUs and JAX versions, demonstrate the error's impact on iterative solvers due to error accumulation, and provide a detailed compiler-level diagnosis pinpointing the bug in XLA's gemm fusion logic; workarounds include lowering the autotune level or restructuring code, and the issue has been cross-posted to the XLA repository for a fix.
   • Number of comments this week: 5

2. [BETTER_ERRORS] c++ pytree code allows unhashable metadata: This issue addresses a problem in the JAX pytree implementation where custom pytree nodes can include unhashable metadata, such as numpy arrays, which leads to a failure during equality checks despite hashing succeeding. The root cause is that the hashing mechanism only hashes pointers to metadata for performance reasons, deferring the actual equality comparison to a later stage where numpy's array equality semantics cause a ValueError, resulting in unclear error messages and crashes in scenarios like JIT caching.

   • The comments discuss potential solutions including adding clearer error messages when unhashable metadata is detected, explain the current hashing and equality checking implementation in C++ that leads to this issue, and consider approaches to detect or prevent these errors earlier without significant performance costs; some progress has been made in documentation and error reporting, and related issues and pull requests are referenced.
   • Number of comments this week: 3

3. [BUG] Performance regression in batched jnp.linalg.solve on CPU from jax/jaxlib 0.9.2 to 0.10.0: This issue reports a significant performance regression in batched small linear solves on CPU when upgrading from jax/jaxlib version 0.9.2 to 0.10.0, with benchmarks showing roughly a 2.6x slowdown. The regression appears linked to changes in the CPU batching strategy for LAPACK kernels, specifically the introduction of parallel batch chunking and thread-pool usage, which negatively impacts performance on large batches of small systems, and is further complicated by differences in OpenBLAS parallelization support.

   • The comments analyze the cause of the regression, identifying that the algorithm for jnp.linalg.solve remains the same but the batching strategy changed to parallel processing in the newer version, which leads to slower performance on certain CPUs. Users confirm the regression on Linux but not on Mac, discuss the impact of OpenBLAS parallelization removal, and note the difficulty of building with Eigen BLAS externally, suggesting that shipping a custom LAPACK/BLAS build might be a future solution.
   • Number of comments this week: 3

4. [BUG] jax.nn.dot_product_attention incompatible with jax.jacobian: This issue reports that the function jax.nn.dot_product_attention is currently incompatible with jax.jacobian, resulting in a TypeError related to improper handling of the vmap axis during reshaping in the backward pass. The user provides a minimal reproducible example and suspects the problem lies in how the batch dimension is managed but is unsure of the correct fix.

   • The comments indicate interest in addressing the issue, with one contributor attempting a fix and noting related bugs in fp8 batchers, leading to separate pull requests that build on each other to resolve similar problems.
   • Number of comments this week: 2

5. [BUG] vmap using in_axes doesn't handle named arguments: This issue reports a problem with the jax.vmap function when using the in_axes argument alongside named (keyword) arguments, which leads to an AssertionError and prevents proper vectorization of functions with keyword parameters. The user provides a minimal example demonstrating that while positional arguments work correctly with vmap and in_axes, passing the same arguments as named parameters causes a failure, highlighting a limitation in handling keyword arguments during vectorization.

   • The comments confirm the issue persists across versions and discuss that the error message is unhelpful and confusing. Users share minimal reproducible examples, express the need for better error handling or explicit errors, and suggest potential workarounds such as wrapping functions or using partial application. There is also interest in extending vmap to support dictionary-style in_axes for keyword arguments, and some attempts to improve the error message are noted.
   • Number of comments this week: 1

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: 6

Summarized Issues:

• GPU and TPU Backend Computation Errors: There are silent miscompilations and crashes occurring in GPU and TPU backends that cause incorrect computations or process termination. The GPU JIT compiler on NVIDIA RTX 5090 misapplies scalar multipliers causing deterministic errors, while the TPU Pallas kernel's block-local vector gather either crashes or requires inefficient payload replication.
• [issues/39336, issues/39366]

• CPU Backend Numerical and Build Issues: The CPU backend experiences both numerical inaccuracies and build failures that do not affect other platforms. The jnp.arctan function returns incorrect values for large arrays in jaxlib 0.10.2, and building JAX 0.11.0 from source on x86_64 fails due to improper header inclusion.
• [issues/39352, issues/39428]

• API Availability and Performance Regressions: Users face limitations due to unavailable stable APIs and performance regressions caused by inefficient code paths. The Pallas pipelining scheduler API is only on the main branch without an official release, and using lax.cond with a no-op branch causes unnecessary array copies and performance drops, with a suggested workaround using lax.switch.
• [issues/39397, issues/39449]

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: 8

Summarized Issues:

• Documentation and Profiling Enhancements: The shard_map documentation was updated to include a pad-and-add pattern for bidirectional collective matrix multiplication, along with code examples and profiling notes. Additionally, there was a request to add profiling tools that provide performance information at a within-kernel granularity level for JAX, aiming to improve performance analysis capabilities.
• issues/20243, issues/21338

• Serialization and Deprecation of Pallas Triton: An attempt was made to enable serialization of Pallas Triton calls as PTX via jax.export to improve stability, but this was closed due to the deprecation of Pallas Triton in favor of Pallas MGPU and jax-triton. This reflects a shift in the underlying tooling and the need to adapt to newer frameworks.
• issues/25196

• Runtime Errors in Parallelism with TensorFlow: A runtime error involving an NCCL internal failure occurred when using JAX parallelism with TensorFlow 2.20.0 on multiple NVIDIA A100 GPUs, causing replicated computation to fail. This highlights challenges in interoperability and stability in multi-GPU environments.
• issues/34918

• Prime Factorization Bug in Logical Mesh Enumeration: The _get_prime_factors function contained a bug where smaller prime factors were dropped if the largest prime factor exceeded the integer square root, leading to incorrect factorization and downstream errors in logical mesh enumeration. This bug affected the correctness of mesh computations relying on prime factorization.
• issues/38286

• Unsupported Input Types Causing Confusing Errors: The functions jax.scipy.linalg.inv and jax.scipy.linalg.cholesky accepted unsupported float16 and bfloat16 inputs without validation, resulting in confusing internal KeyError exceptions deep within LAPACK's dtype dispatch. This issue pointed to the need for clearer upfront user-facing error messages for unsupported dtypes.
• issues/38825

• Incorrect Results Due to Reduction Fusion Bug: In JAX 0.10.2, using jax.jit on a vmapped matrix-vector product followed by a reduction on CPU returned incorrect results because of a YNNPACK reduction fusion problem, which caused discrepancies compared to eager execution. This bug was fixed in version 0.11.0, restoring correct behavior.
• issues/39103

• Regression Causing Sharding Mismatch on TPU: A regression in JAX 0.11.0 caused reshaping an array with a size-1 explicit mesh axis to drop that axis, leading to sharding mismatch errors during concatenation on TPU devices. This issue did not occur in JAX 0.10.2 and affected multi-device array operations.
• issues/39309

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: 17

Key Open Pull Requests

1. Prioritize high-bandwidth axes in device mesh creation for TPU v7: This pull request improves device mesh creation for TPU v7 by adding a physical_axis_priority parameter that ensures high network-intensity logical axes are preferentially assigned to higher-bandwidth physical axes—specifically prioritizing the core axis over x, y, and z axes in the 4D physical mesh topology—to optimize bandwidth utilization during communication-intensive operations like FSDP.

• URL: pull/39340

• Associated Commits: ef56b, 24a70, 10822, 4c464, 45d3c

2. [ROCm] Fix invalid parallel local jobs execution: This pull request addresses the issue of invalid parallel execution of local test jobs in the ROCm environment by limiting the number of concurrent local test jobs to one, due to the use of a single GPU GitHub Actions runner.

• URL: pull/39419

• Associated Commits: 5e92e, 816ba

3. Fix cudnn fp8 attention batching rules: This pull request fixes the batching rules for the cudnn fp8 attention primitives in JAX by normalizing tensor operands through broadcasting the vmap axis onto unbatched operands, correctly handling output shapes and labeling outputs to ensure compatibility with vmap operations, thereby resolving crashes and errors encountered in fp8 attention under various vmap scenarios.

• URL: pull/39452

• Associated Commits: 81827, ab3ca

Other Open Pull Requests

• Documentation fixes and updates: Several pull requests improve the accuracy and clarity of documentation across the JAX project. These include fixing grammatical and mathematical typos in jax.scipy.special, updating docstrings for lax.conv functions to clarify default dimension orders, reflecting changes from issue #39198, and adding gradient printing to the autodiff tutorial example.
  • pull/39361, pull/39411, pull/39447, pull/39405

• Bug fixes and correctness improvements: Multiple pull requests address correctness issues and bugs in the codebase. These include fixing a heap corruption issue caused by OpenBLAS thread oversubscription, correcting the missing increment of an iterator in FlattenOneLevelImpl to handle namedtuples properly, fixing the stop_gradient transpose for batched conditional operations, and correcting cudnn fused attention batching rules to handle unbatched operands correctly.
  • pull/39383, pull/39448, pull/39453, pull/39451

• Error handling improvements: Some pull requests enhance error messages and detection mechanisms to improve user experience. This includes detecting stale tracers in custom_jvp calls and raising a clear UnexpectedTracerError, as well as improving the vmap in_axes arity check error message by explicitly mentioning keyword arguments and suggesting fixes.
  • pull/39386, pull/39450

• New features and functionality extensions: A few pull requests add new capabilities or extend existing functionality. These include introducing a faster, loop-free approximate Poisson random number generator method supporting float outputs and derivatives, extending pl.empty_ref_like to accept Ref inputs with support for AbstractRef abstract values, and adding tests to verify custom_vjp functionality with hi types in the hijax module.
  • pull/39353, pull/39354, pull/39441

• Testing improvements: Some pull requests focus on increasing test coverage and ensuring correctness through additional tests. This includes adding new unit tests for the Poisson RNG and the function jax.ds(), as well as tests for custom_vjp with hi types.
  • pull/39353, pull/39461, pull/39441

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: 28

Key Closed Pull Requests

1. dtypes: fix missing canonical dtype name in ExplicitX64Mode.ERROR message: This pull request fixes the error message in the _maybe_canonicalize_explicit_dtype function's ERROR mode by adding a missing placeholder to correctly display the canonical dtype name when a 64-bit dtype is requested with enable_x64=False, aligning it with the WARN mode message, and includes a new test to ensure this correction is properly validated.

• URL: pull/38033

• Associated Commits: 001d8, 0f290, 5f1f7

• Associated Commits: 001d8, 0f290, 5f1f7

2. Fix mesh prime factorization with large remainders: This pull request addresses a bug in the mesh prime factorization function by preserving previously found factors when the remaining quotient is a larger prime, thereby fixing incomplete factorization for certain axis sizes like 10 and 22, and includes regression tests to validate the correction.

• URL: pull/38574

• Associated Commits: ea73f, 02500

• Associated Commits: ea73f, 02500

3. Parameterize sampler tests by whether or not to jit: This pull request refactors the sampler tests in random_lax_test.py to parameterize them by whether or not to use JIT compilation, aiming to reduce expensive compilation times and improve test runtime efficiency when subsampling test cases.

• URL: pull/38972

• Associated Commits: e1695, 5eafc

• Associated Commits: e1695, 5eafc

Other Closed Pull Requests

• _get_prime_factors bug fixes: Multiple pull requests fix a critical bug in the _get_prime_factors function where the use of a for…else loop without a break caused smaller prime factors to be dropped when the largest prime factor was large, leading to incorrect factorizations and breaking device mesh construction with allow_split_physical_axes=True. These fixes remove the for…else construct, append the leftover prime factor after the loop, and add regression tests to ensure correctness and maintain the invariant prod(factors) == x.
• [pull/38287, pull/38466, pull/39300, pull/39324]

• LAPACK dtype error handling improvements: Two pull requests improve error handling for unsupported LAPACK data types in CPU linear algebra operations by adding explicit dtype checks that raise clear and actionable NotImplementedError messages instead of confusing internal errors caused by dtype-to-prefix lookup failures. These changes refine the error messages to avoid confusion while maintaining compatibility with existing tests.
• [pull/39295, pull/39381]

• remat3 feature fixes and enhancements: Several pull requests address failures, regressions, and enhancements in the remat3 feature by fixing constant attachment bugs, supporting mutable arrays, enabling nested rematting, forwarding reference residuals through remat, and improving effect handling and test stability. These changes also include adding a dedicated remat rule for RematTraced, extending optimization_barrier to handle symbolic zeros and refs, and skipping a regressed key-reuse-checking test to maintain test suite health.
• [pull/39429, pull/39432, pull/39438, pull/39439, pull/39444]

• random and testing improvements: Pull requests simplify tests for the random.gamma function and remove redundant just-in-time compilations in sampler tests to reduce overhead and improve test execution time. Additionally, version pinning for Hypothesis and mpmath testing libraries is removed or adjusted, with fixes to underlying test failures and suppression of deprecation warnings to stabilize the test environment.
• [pull/38761, pull/39356, pull/39358, pull/39401, pull/39462]

• ROC remote build executor and backend updates: Two pull requests reduce the parallelism of ROCm remote build executor jobs to lower system load and update the ROCm backend to reference the latest XLA version, ensuring compatibility and improved performance.
• [pull/39400, pull/39424]

• PyTreeDef error reporting enhancement: One pull request improves error reporting in the PyTreeDef equality operator by including the representation of the specific registered pytree type that caused metadata comparison failures, making it clearer which pytree type triggered the error when node metadata is unhashable or uncomparable.
• [pull/39427]

• custom_jvp3 feature introduction: A pull request introduces a new custom_jvp3 feature where the user's rule is defined by the jvp method, with linear and reverse-mode derivatives automatically derived, and includes vmap batchedness inference from the primal function alone, improving error handling for outputs with unexpected batching.
• [pull/39440]

• Jaxpr class attribute removal: One pull request removes the Jaxpr._num_consts attribute from the Jaxpr class to streamline the codebase.
• [pull/39323]

• shard_map matrix multiplication refactor: A pull request refactors the shard_map bidi collective matrix multiplication example by replacing two half-size matrix multiplies per ring step with a single full-size multiply using a pad-and-add pattern, improving efficiency while maintaining correctness.
• [pull/39413]

• RBE pool switch proposal: One pull request proposes switching the JAX project to use the mi350 RBE pool instead of the mi300 pool due to issues with the latter, but this change was not merged.
• [pull/39420]

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.