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 without introducing major changes. This release highlights a focus on incremental fixes and refinements.

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. [BETTER_ERRORS] Transpose rule for 'custom_vjp_call_jaxpr' not implemented for grad(custom_jvp(custom_vjp)): This issue reports a NotImplementedError encountered when attempting to compute the gradient of a function defined with a custom_jvp that internally calls a custom_vjp, due to the absence of a transpose rule for the custom_vjp_call_jaxpr primitive in JAX's reverse-mode autodiff. The user suggests that this limitation arises because the custom_vjp_call primitive appears inside the tangent computation of the custom_jvp and lacks a registered transpose rule, proposing possible solutions including registering a transpose rule, routing through existing machinery, or improving the error message to guide users toward appropriate alternatives.

   • The comments confirm the issue reproduces on newer JAX versions and provide a detailed root cause analysis explaining why the transpose rule is missing. They discuss potential fixes and converge on the idea of improving the error message to direct users to use hijax primitives like VJPHiPrimitive for defining transpose rules, with plans to document this approach and update the error messaging accordingly.
   • Number of comments this week: 4

2. [BUG] [Pallas:TPU:SparseCore] Support local offset selection after indirect DMA block gather: This issue concerns the inability to perform local offset selection after an indirect DMA block gather operation in the SparseCore pipeline on TPU v6e, which currently requires materializing the entire block to high-bandwidth memory (HBM), leading to significant performance costs. The user requests support for block-local dynamic offset selection or a scalar/vector-output indirect DMA form to avoid this costly workaround and improve efficiency in handling blocked embedding tables or compact primary-key indexes.

   • The comments describe testing a proposed primitive that lowers correctly but initially fails at runtime due to unsupported operations; a workaround disabling layout passes enables successful execution and exact results, with ongoing work to improve vector layout pass support to remove the need for this workaround.
   • Number of comments this week: 4

3. [BUG] jnp.nanquantile fails on empty reduction axes: This issue reports that the function jnp.nanquantile in JAX raises an indexing or gather error when the reduction axis is empty, whereas the equivalent NumPy function returns NaNs with the expected reduced shape. The user provides a minimal reproducible example and notes that the failure occurs under various conditions including different axis specifications, keepdims settings, and both eager and JIT-compiled execution.

   • The comments reveal that two separate pull requests have been opened to fix the issue by aligning JAX's behavior with NumPy's, with requests for maintainers to add review labels to these PRs.
   • Number of comments this week: 3

4. [ENHANCEMENT] [CONTRIBUTIONS WELCOME] scipy.linalg.solve_banded and scipy.linalg.ldl: This issue discusses the potential addition of two linear algebra functions, scipy.linalg.solve_banded and scipy.linalg.ldl, to JAX, highlighting their importance in optimization tasks and the relative ease of wrapping existing LAPACK routines. The conversation focuses on the challenges of GPU support, particularly for solve_banded, and explores possible implementation strategies, while noting that ldl is less blocked and could be integrated sooner with community contributions.

   • The comments reveal interest in these functions for optimization, provide guidance on extending JAX with new linear algebra primitives, and discuss the technical difficulties of GPU support for banded solvers, with suggestions to prioritize ldl integration first and consider fallback options for solve_banded; community members express willingness to contribute implementations and discuss potential external libraries for GPU acceleration.
   • Number of comments this week: 2

5. [BUG] CPU performance regression: This issue reports significant CPU performance regressions in JAX across multiple versions, with runtimes worsening notably after version 0.4.31, particularly following the introduction of the thunk runtime. The user provides a reproducible example and timing data, expressing frustration with CPU performance and advising users to remain on version 0.4.31 until improvements are made.

   • Comments reveal that the regression worsened in versions 0.10.2 and 0.11.0, but a workaround using the --xla_cpu_experimental_ynn_fusion_type=-reduce flag significantly improved performance by addressing a fusion bug in XLA; however, this fix only partially mitigates the regression on some hardware.
   • 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: 13

Summarized Issues:

• Reduction and indexing errors in JAX nanquantile and multi-output fusion miscompilation: JAX's jnp.nanquantile raises errors when the reduction axis is empty, unlike NumPy which returns NaNs with the expected shape, showing inconsistent behavior across axis specifications and execution modes. Additionally, a silent miscompilation in XLA's Triton multi-output fusion on NVIDIA Blackwell GPUs causes large numerical errors without crashes, triggered by complex fusion patterns involving nested concatenates and multiple reductions.
• [issues/39468, issues/39486]

• GPU compile time and multi-GPU test failures on NVIDIA hardware: There is a significant increase in GPU compile time—about 190x longer—when performing trailing reductions after scatter-add chains, while CPU compile time remains unaffected, raising questions about profiling and optimization. Multi-GPU tests fail on NVIDIA B200 and GB200 devices due to StableHLO parser errors with nested replica groups and CUDA runtime errors preventing Blackwell GPU initialization under CUDA 13, despite correct environment setup.
• [issues/39483, issues/39659, issues/39660]

• SparseCore and TPU memory and DMA issues: Several issues affect SparseCore and TPU memory operations, including a DMA round trip returning zero for 4-byte payloads on Pallas TPU, lack of TPU v4 CMEM scratch memory allocation support causing runtime errors, and failure of SparseCore consumer kernels on TPU v6e due to device-side halt errors when consuming lookup buffers from ordinary JAX computations. These problems highlight challenges in memory handling and kernel interoperability on TPU hardware.
• [issues/39505, issues/39615, issues/39662]

• SparseCore DMA and data type support limitations on TPU: SparseCore indirect DMA operations on TPU face limitations such as lack of support for bf16 data types causing runtime errors, and the inability to perform local offset selection after indirect DMA block gather on TPU v6e, forcing inefficient full block materialization. There is also a request for SparseCore-to-TensorCore Remote DMA support on Cloud TPU to enable currently skipped Pallas kernel features.
• [issues/39518, issues/39577, issues/39621]

• Symmetric tensor storage format request: A request was made for a memory-efficient storage format for symmetric tensors of arbitrary order that stores only unique components while supporting indexing, matrix-vector products, automatic differentiation, and compatibility with JAX transformations like jit, vmap, and grad. This would optimize memory usage and computation for symmetric tensor operations in JAX.
• [issues/39503]

• Incorrect results from donate_argnums on CUDA GPUs: Using donate_argnums in JAX on CUDA GPUs produces incorrect results depending on the output order when returning tuples with multiple items, an issue not observed on CPUs or with fewer tuple items. This indicates a GPU-specific bug affecting function output correctness in certain tuple return scenarios.
• [issues/39625]

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

Summarized Issues:

• Gradient computation errors with array outputs: This issue addresses a TypeError in JAX when computing gradients of functions returning arrays shaped (1,) instead of scalars, questioning whether such arrays should be treated as scalars. It also discusses the challenges of using .item() or slicing to convert outputs and seeks advice on proper handling of these gradient computations.
• issues/2303

• Improving error messages for indexing issues: This topic covers the enhancement of IndexError messages to include more specific information about the cause, such as the JAX transformation and user source line responsible for the error. The goal is to provide clearer diagnostics for debugging indexing problems.
• issues/7222

• Inefficiencies and incompatibilities in dot_product_attention: Multiple issues highlight inefficiencies in the cudnn implementation of dot_product_attention due to unnecessary dbias computations when masks are used, and incompatibility of jax.nn.dot_product_attention with jax.jacobian caused by improper handling of vmap axes during reshaping in the backward pass. These problems lead to increased runtime and TypeErrors during gradient computations.
• issues/34685, issues/38495

• Configuration and environment variable handling bugs: This issue reports that setting jax.config.update("jax_compiler_enable_remat_pass", False) does not properly respect existing XLA_FLAGS environment variables, resulting in incorrect compilation options being applied in JAX.
• issues/37391

• Key assignment bug in flatten_one_level_with_keys: A bug in JAX 0.10.2 causes the function flatten_one_level_with_keys to assign the first namedtuple field key to every child element instead of distinct keys, leading to repeated keys in the flattened output.
• issues/39297

• Performance improvements for random sampling: This issue proposes replacing the loop-based implementation of jax.random.poisson with a faster approximate arithmetic-only sampler to significantly increase throughput and reduce latency on GPUs while maintaining minimal approximation error.
• issues/39312

• API availability and release timeline inquiries: This issue requests information about the timeline for including the Pallas pipelining scheduler API, currently only on the main branch, in an official released wheel to enable its supported public use without relying on main builds.
• issues/39397

• Batching rule output mismatch bug: The dce_sink batching rule incorrectly emits an output for a primitive that should produce none, causing a mismatch error inside a vmapped checkify(while_loop) construct. A fix is proposed to make the batching rule emit no outputs to maintain consistency.
• issues/39501

• DMA API extension for static permutation support: This issue requests support for transferring data between HBM and VMEM with static permutation of non-tiled axes during DMA operations, proposing an API extension to pltpu.make_async_copy() to handle such permutations directly and avoid separate layout conversions.
• issues/39532

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

Key Open Pull Requests

1. [jax.numpy] Fix nanquantile on empty reduction axes: This pull request fixes the jnp.nanquantile function in JAX to correctly handle cases where the reduction axis is empty by returning a NaN-filled array matching NumPy's behavior, preventing indexing errors, and includes regression tests covering various axis configurations, quantile values, and execution modes.

• URL: pull/39475

• Associated Commits: 2e7f0, d693d, 07db4, e243a, dbe1a, f5e1b

2. Point custom_vjp transpose error to hijax VJPHiPrimitive; document hi-primitive transpose: This pull request improves error handling and documentation for the transpose rule of custom_vjp functions nested inside custom_jvp tangent rules by registering a targeted error message directing users to use the VJPHiPrimitive hijax primitive with its own transpose method, adding detailed documentation with examples for hi-primitive transposes, and including regression tests to ensure clarity and correctness in these scenarios.

• URL: pull/39589

• Associated Commits: 6eb7d, d1402, 0348f

3. Removed deprecated pl.pallas_call from mgpu_torch_test: This pull request removes the deprecated pl.pallas_call from the mgpu_torch_test module, improving compatibility and following updates from a related prior pull request.

• URL: pull/39678

• Associated Commits: 8a700, b0b5d

Other Open Pull Requests

• GPU and ROCm CI Improvements: Multiple pull requests enhance GPU-related functionality and ROCm CI workflows. These include recoding gpu_info using jaxlib._gpu_spec, adding support for Rubin GPU and unnamed NVIDIA devices, fixing kernel verification errors, and improving ROCm container teardown by adding the --init flag and cleanup steps to prevent hangs.
• [pull/39469, pull/39473, pull/39546, pull/39547, pull/39593, pull/39601, pull/39619]

• OneAPI Backend Enhancements: Several pull requests add foundational support and callback integration for the OneAPI backend. They introduce GPU Python callback infrastructure, register oneMKL-based solver kernels including Cholesky, fix startup failures in OneAPI workflows, and improve GPU runtime with asynchronous memory operations.
• [pull/39601, pull/39632, pull/39651]

• Bug Fixes in GPU Kernels and Attention Mechanisms: Pull requests address critical bugs in GPU kernel compilation and cuDNN sliding window dot product attention. Fixes include correcting indexing errors, ensuring proper LLVM NVPTX target registration, and resolving attention masking issues with comprehensive test coverage.
• [pull/39561, pull/39579, pull/39639]

• Documentation and Test Improvements: Several pull requests improve documentation and testing. These include clarifying buffer_callback docstrings with GPU-compatible examples and regression tests, fixing typos in pipelining markdown, updating device memory profiling commands, adding skip decorators for TPU tests, and adding regression tests for XLA dynamic slice fusion.
• [pull/39476, pull/39478, pull/39576, pull/39606, pull/39490]

• Code Cleanup and Refactoring: Some pull requests focus on code simplification and removal of legacy workflows. This includes removing call_p and closed_call_p functions in favor of eval_jaxpr_p, deleting the legacy ROCm GPU Post-Merge Check workflow, and fixing unpickling crashes in jaxlib by adding validation and runtime checks.
• [pull/39593, pull/39632, pull/39658]

• Numerical and API Behavior Fixes: Pull requests fix subtle numerical behavior and API issues. These include preserving the sign of zero in complex sqrt results to match NumPy, correcting shape indexing in matmul_small, and adding proper rejection of invalid array-valued indices in SparseCore kernels.
• [pull/39579, pull/39640, pull/39547]

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

Key Closed Pull Requests

1. [ROCm] Add TheRock latest to nightly and continuous CI workflows: This pull request updates the ROCm continuous integration workflows by adding the latest TheRock release to the nightly and continuous CI matrices for build, pytest, and bazel jobs, aligning the configuration with that of the ROCm/jax amd-main branch and ensuring comprehensive testing across multiple TheRock release tracks.

• URL: pull/38810

• Associated Commits: c8473, ab79d, 58614, 15b1e, 8a193

2. [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, ensuring reliable and consistent artifact generation.

• URL: pull/39282

• Associated Commits: c8ca6, 3367e, e7d64

3. Fix cudnn fp8 attention batching rules: This pull request fixes the batching rules for cudnn fp8 attention 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 batching scenarios.

• URL: pull/39452

• Associated Commits: 81827, ab3ca

Other Closed Pull Requests

• Hermetic build improvements: Several pull requests focus on improving the hermetic build process for JAX, including using the hermetic clang compiler from XLA toolchains, switching ROCm builds to hermetic LLVM, and adding runfiles data files and ROCm hermetic dependencies to plugin and PJRT wheels. These changes aim to standardize, isolate, and enable proper hermetic Bazel test execution for more consistent builds.
  • pull/36283, pull/37071, pull/39220

• Namedtuple and pytree flattening fixes: Two pull requests address issues with namedtuple keys in one-level pytree flattening by fixing bugs in the flatten_one_level_with_keys method and adding tests to verify the fix. One of these was not merged, but both contribute to handling namedtuples correctly in pytrees.
  • pull/39306, pull/39448

• Error message improvements: Multiple pull requests enhance error messages for better user guidance, including clarifying vmap in_axes arity checks with keyword arguments, improving error messages for non-static slice indices in NumPy-style slices, and providing clearer instructions when jax.grad rejects non-scalar outputs. These improvements help users understand and fix common errors more easily.
  • pull/39450, pull/39502, pull/39557

• Remat and optimization barrier adjustments: Several pull requests modify the remat3 implementation and optimization barriers to prevent numerical breakages and improve compatibility. They include moving the optimization_barrier to match remat2, introducing a flag to exclude outgrad from the CSE barrier, and fixing internal dead code elimination for RematTraced. These changes address downstream failures and improve remat behavior.
  • pull/39484, pull/39512, pull/39513

• cuDNN attention and packed-attention fixes: Pull requests fix issues in cuDNN attention implementations by skipping dbias gradient computation for constant boolean masks and correcting kv_offsets calculation for grouped-query and multi-query attention. These fixes prevent unnecessary gradients and silent errors in attention mechanisms.
  • pull/39508, pull/39520

• Testing and documentation updates: Some pull requests add tests for custom_vjp functionality with hi types, unit tests for jax.ds(), and update documentation to reflect recent changes and improve readability. These efforts enhance code coverage and user understanding.
  • pull/39441, pull/39447, pull/39461, pull/39493

• Build and CI workflow enhancements: Pull requests introduce a dedicated oneAPI continuous integration workflow to build hermetic JAX wheel artifacts and modify the build process to compile ROCm plugin artifacts on CPU runners, freeing GPU resources for tests requiring actual GPUs. These changes improve build efficiency and resource utilization.
  • pull/39470, pull/39521

• Serialization and versioning improvements: One pull request introduces JAX serialization version 12 by adding the JAX version used during export serialization and a flag to bypass deserialization version errors temporarily. This update enhances serialization robustness and backward compatibility.
  • pull/39494

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.