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 an incremental update to v0.9.0 that incorporates patches from four specific pull requests, focusing on targeted improvements and bug fixes without introducing 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. Pallas-TPU: arith.shrui on int8 vectors fails Mosaic legalization (blocks in-kernel nibble unpack): This issue describes a failure in the Mosaic legalization process when performing a logical right-shift operation on int8 vectors within a Pallas-TPU kernel, specifically during the standard nibble-unpack operation for 4-bit packed weights. The problem arises because the Mosaic backend does not support the arith.shrui operation on int8-typed vectors, although casting to int32 before shifting serves as a successful workaround.

   • The comments discuss potential solutions including adding compiler emulation that mimics casting to 32-bit before shifting and back, suggest a more efficient workaround involving bit manipulation and casting, and debate whether compiler emulation is necessary given the efficiency of the casting approach; a related pull request is also referenced.
   • Number of comments this week: 4

2. [BUG] jnp.bfloat16's smallest_subnormal truncated to zero when converted to jnp.float64 and then back to jnp.bfloat16.: This issue reports that when converting the smallest subnormal value of jnp.bfloat16 to jnp.float64 and then back to jnp.bfloat16, the value is truncated to zero, which does not happen when converting from other types like jnp.float32. The user finds this behavior unexpected because the smallest subnormal is representable in jnp.float64, and the discrepancy arises due to how flush-to-zero semantics are applied differently in these conversions.

   • The comments explain that XLA generally uses flush-to-zero semantics for subnormal values, which accounts for the truncation to zero; however, the user questions why this behavior differs between float32 and float64 conversions. It is clarified that conversions between float32 and bfloat16 are simpler and do not involve value-dependent logic, whereas conversions involving float64 are more complex and thus subject to flush-to-zero rules.
   • Number of comments this week: 3

3. [BUG] [XLA] [TPU] Silent wrong gradient: eager (non-jit) VJP of x[:, idx, :] gather zeroes batch rows on TPU (regression in libtpu 0.0.42): This issue describes a regression in the TPU backend of JAX where the eager (non-jit) vector-Jacobian product (VJP) of a gather operation silently produces incorrect gradients by zeroing out half of the batch rows, while the same operation under jit returns correct results. The problem appears in specific versions of JAX and libtpu and depends on the input dimensions, causing a discrepancy that affects gradient computations on TPU devices.

   • The comments acknowledge the issue as a likely XLA:TPU bug and mention that an internal report has been filed to track the problem.
   • Number of comments this week: 2

4. [BUG] GPU scatter-add (x.at[idx].add(y)) into a small array up to ~2000x slower in jaxlib 0.10.2 vs 0.10.1: This issue reports a significant performance regression in the GPU scatter-add operation (x.at[idx].add(y)) when upgrading from jaxlib version 0.10.1 to 0.10.2, with slowdowns up to approximately 2000 times for small output arrays due to increased index contention. The user provides a detailed reproducible benchmark script and timing results on an NVIDIA RTX A4000, showing that while large output sizes recover performance, small sizes suffer a drastic slowdown, and suspects the root cause lies in changes to XLA's backend code generation rather than the high-level JAX code itself.

   • Comments confirm similar experiences with the regression, and one user shares HLO IR output showing identical high-level compiled code between versions, reinforcing the hypothesis that the slowdown originates from backend codegen changes in XLA rather than JAX frontend changes.
   • Number of comments this week: 1

5. [ENHANCEMENT] Feature request: add differentiation rule for jax.scipy.linalg.schur (NotImplementedError on jax.grad): This issue reports that the function jax.scipy.linalg.schur currently raises a NotImplementedError when used with jax.grad, preventing differentiation through Schur decomposition in JAX. The user requests either the implementation of a correct differentiation rule or a clearer error message that explains the limitation and suggests alternative differentiable matrix operations.

   • The comments express willingness to contribute a solution, proposing to first add a user-friendly error handler to clarify the limitation and then work towards implementing a full reverse-mode autodiff rule for the Schur decomposition, requesting assignment to proceed with this phased approach.
   • 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: 18

Summarized Issues:

• Multi-GPU and TPU Backend Bugs: Several issues describe bugs and regressions in multi-GPU and TPU backends, including silent incorrect results from multi-output fusion aliasing on four CUDA GPUs, incorrect eager VJP gradients on TPU gather operations, and segmentation faults when initializing GPU/TPU JAX on CPU-only machines. These problems cause data corruption, silent failures, or incorrect computations that complicate multi-device usage and debugging.
• issues/39100, issues/39025, issues/39101

• Pallas TPU Kernel and Compiler Limitations: Multiple issues highlight limitations and missing documentation in Pallas TPU kernels, such as legalization failures of arith.shrui on int8 vectors, verification errors with vector.extract_strided_slice due to unsupported stride constraints, and the need to document uint32 source and block-tiling constraints for pltpu.unpack_elementwise. These gaps cause compilation failures or late lowering errors that require workarounds or improved error messages.
• issues/38986, issues/38987, issues/39058

• Numerical and API Inconsistencies in JAX NumPy and SciPy: Several issues report discrepancies between JAX and standard NumPy/SciPy behavior, including jax.numpy.heaviside returning the second input on NaN instead of propagating NaN, jax.numpy.sqrt mishandling complex inputs with negative zero imaginary parts, and jax.scipy.special.betaln returning NaN instead of infinite values for infinite inputs. These inconsistencies lead to unexpected numerical results and reduced compatibility with NumPy standards.
• issues/39109, issues/39110, issues/39106

• Memory and Performance Issues: There are reports of performance degradation and memory management problems, including a 20x slowdown of large batched FFTs when combined with adjacent operations in a single jit, and the pjit cache reaching capacity and permanently disabling fastpath retention, causing slower dispatch for all subsequent jitted functions. These issues impact runtime efficiency and scalability.
• issues/39081, issues/39078

• Data Conversion and Callback Behavior: One issue highlights a problem where converting the smallest subnormal jnp.bfloat16 value to float64 and back truncates to zero, unlike other conversions that preserve the value. Another issue discusses confusion around jax.experimental.io_callback with sharded multi-GPU arrays, where callbacks receive data resharded to a single device rather than per-device shards, calling for better documentation or interface improvements.
• issues/38995, issues/38968

• Build and Compatibility Errors: There is a reported TypeError when building jaxlib 0.10.2 with Python 3.10 due to an unsupported 'required' keyword in argparse's _SubParsersAction.__init__(). This build failure prevents successful installation and usage on newer Python versions.
• issues/38969

• Incorrect Results from JIT and Fusion on CPU: An issue reports that jax.jit on a vmapped matrix-vector product followed by reduction on CPU returns incorrect results in jaxlib 0.10.2, likely due to a YNNPACK fusion or tiling bug triggered by the reduction, differing from eager execution and previous versions.
• issues/39103

• Request for Expanded HLO Modification Support: There is a feature request to expand the jax._src.lib library to support more callable functions for modifying HLO, similar to tensorflow.compiler.xla, to improve flexibility in HLO transformations.
• issues/39023

• TPU Virtual Memory Limit Documentation: One issue requests documentation and support for increasing the default 64MB virtual memory limit for fused dequant kernels on TPU7x devices, as these kernels slightly exceed the limit and require a higher threshold to compile and run correctly.
• issues/38988

• JIT Compilation and Padding Miscompilation: A miscompilation bug in the XLA CPU backend causes zero-padding applied via jnp.pad to be read back as uninitialized or garbage values when JIT-compiled with nested lax.scan and dynamic_slice, resulting in flaky and incorrect results not seen in eager execution.
• issues/39120

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

Summarized Issues:

• Performance and Compilation Overheads: Several issues highlight unexpected slowdowns and overheads in JAX, including slower runtime of lax.scan compared to recursion, slower Jacobian computations with jit(jacfwd(fun)) versus Autograd, and the large initial compilation delay for JIT-compiled Jacobian functions. These problems affect inference speed and user experience, with some exploring static argument handling and caching to optimize performance.
• [issues/2251, issues/2302, issues/4495, issues/4937, issues/39089, issues/26812]

• Control Flow and Conditional Evaluation Challenges: Users face difficulties with lax.cond when auxiliary outputs differ structurally between branches, and crashes occur when combining custom_jvp with nondiff_argnums inside lax.cond. Additionally, using lax.cond within lax.reduce caused shape mismatch errors, indicating challenges in managing control flow and tracing in JAX.
• [issues/2660, issues/9374, issues/12615]

• Array and Data Structure Limitations: JAX's inability to create arrays containing None or heterogeneous types limits list-like semantics, and boolean indexing expressions require refactoring for JIT compatibility due to static shape requirements. Moreover, dictionary key order is not preserved in tree_util flattening/unflattening, causing unexpected behavior compared to Python's insertion order guarantees.
• [issues/1580, issues/2765, issues/4085]

• Gradient and Hessian Computation Issues: Users report challenges in efficiently computing gradients and Hessians, such as obtaining individual Hessians per timestep in scanned functions and implementing custom JVP/VJP rules for multi-linear functions to improve performance. These issues highlight the complexity of advanced differentiation patterns in JAX.
• [issues/3353, issues/3366]

• Vectorization and Mapping Anomalies: Oscillatory behavior and internal assertion failures occur when using vmap combined with lax.scan or with conflicting axis names in shard_map. These problems cause inconsistent outputs and opaque errors, complicating parallelization and vectorized computations.
• [issues/3814, issues/15905]

• Static Value and Tracer Handling: Nested JIT compilation can unexpectedly convert static values into tracers, causing errors and surprising behavior. This indicates subtle issues in how JAX manages static versus dynamic values during compilation and execution.
• [issues/14666]

• Scatter and Sharded Array Bugs: Scatter operations on sharded matrices cause out-of-memory errors or incorrect updates due to all-gather operations or partial index updates. These bugs affect distributed computations and have been partially addressed but remain challenging for large arrays.
• [issues/23052]

• Random Sampling Discrepancies: The jax.random.choice function with replacement can sample zero-probability indices for large inputs with low average probabilities, differing from NumPy's behavior and potentially causing unexpected sampling results.
• [issues/25498]

• Error Handling and Crash Fixes: Several issues report crashes or assertion failures due to internal errors, such as integer tangent promotion in custom_jvp, invalid attribute access in tree utilities, and type errors when writing to global mutable arrays inside jax.grad functions. Fixes have improved error messages and prevented crashes.
• [issues/16000, issues/25659, issues/26361]

• Documentation and Usability Improvements: Requests for better documentation include clarifying the differences between "high" and "low" jaxprs and improving error messages for unsupported operations like jax.nn.scaled_matmul on CPU. These aim to enhance user understanding and reduce confusion.
• [issues/31892, issues/38813]

• Hardware and Device Support Issues: Users encounter problems such as "No GPU/TPU found" errors resolved by upgrades but still experience slow performance, and requests for programmatic hardware memory monitoring to aid benchmarking and batch size tuning, especially on TPUs.
• [issues/21240, issues/8096]

• Library and Ecosystem Management: Debates arise over updating the JAX README to add or remove neural network libraries, reflecting concerns about stability, maintenance, and fairness in listing downstream projects for new users.
• [issues/20745]

• Mathematical Function Accuracy Bugs: Several special functions in jax.scipy.special return NaN instead of mathematically expected finite values when given infinite arguments, indicating correctness issues in edge cases for beta and incomplete gamma functions.
• [issues/39105, issues/39107, issues/39108]

• Security and Workflow Hardening: A backport of workflow permission restrictions was applied to multiple release branches to maintain minimal security standards and prevent excessive permissions in CI workflows.
• [issues/38835]

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. Remove Jaxpr._num_consts: This pull request removes the attribute Jaxpr._num_consts from the Jaxpr class and includes multiple subsequent fixes and updates, such as modifying eval_jaxpr to use jaxpr consts instead.

• URL: pull/39049

• Associated Commits: dc679, 74b69, 4b33a, e29c8, e4238, 3a4f6, b5a2f, 68f4f, 00ee6

2. docs(nn): describe scaled_matmul platform behavior accurately: This pull request updates the scaled_matmul function's docstring to accurately describe its platform support by clarifying that while GPU backends use a native fused fast path, other backends lower through the xla.scaled_dot composite with fallback behavior, correcting the previous inaccurate claim that it was GPU-only and always raised NotImplementedError on TPU or CPU.

• URL: pull/38954

• Associated Commits: 11734, a5a1a, 93f83

3. fix(lax): add gradient rules for betainc w.r.t. a and b: This pull request adds support for gradients of jax.lax.betainc(a, b, x) with respect to the shape parameters a and b by implementing a term-by-term differentiation of the underlying hypergeometric series, fixing prior errors that raised exceptions when differentiating these parameters, improving convergence criteria to handle integer values correctly, and verifying accuracy against high-precision references, thereby enabling use of betainc in training loops with learnable shape parameters.

• URL: pull/38977

• Associated Commits: c8cc3, 48a06

Other Open Pull Requests

• Documentation enhancements and restructuring: Multiple pull requests add new documentation sections and update existing ones, including a 501 systems topics section covering advanced JAX features, incorporation of the CuTe DSL tutorial into the 401 section, addition of an errors reference to 101, updates to cross-links in 201, and a new landing page organizing documentation with a clear table of contents. These changes improve user guidance on topics like external callbacks, Pallas orientation, and FFI tutorials, enhancing overall documentation clarity and navigation.
  • pull/39095, pull/39007, pull/39088

• Performance and test improvements: Refactoring of sampler tests in random_lax_test.py introduces parameterization by JIT compilation usage, significantly reducing expensive compilation times and improving runtime efficiency during test subsampling.
  • pull/38972

• Pallas-TPU kernel workarounds and fixes: Several pull requests address issues with Pallas-TPU fused kernels, including an int32-cast workaround for the arith.shrui legalization gap on Mosaic backend for int8 vectors, a workaround for vector.extract_strided_slice verification errors by changing weight layout and activation permutation, and increasing VMEM limits to fix out-of-memory errors during matmul operations, all verified on TPU hardware.
  • pull/38990, pull/38991, pull/38992

• GPU device info and compatibility fixes: Fixes to GPU info functionality ensure correct retrieval of compute capability and architecture names from real devices, and a ROCm-specific fix falls back on gfx ISA for Triton target selection to handle unknown device marketing names, making GPU version optional to support future architectures.
  • pull/39000, pull/39038

• Build and argument parsing robustness: A compatibility fix in the build script modifies how the required flag is set for add_subparsers in argparse, changing it from a keyword argument to an attribute to avoid TypeErrors with third-party argparse backports, maintaining subcommand enforcement.
  • pull/39001

• PJRT GPU plugin update: The PJRT GPU plugin in jaxlib is updated to forward the custom call traits argument introduced in extension version 3, enabling support for this feature while maintaining backward compatibility through conditional compilation.
  • pull/39003

• PRNG and oneAPI integration: A SYCL/oneAPI implementation of the ThreeFry2x32 PRNG GPU kernel is added for Intel GPUs, including registration, lowering, build, packaging, and extended testing to ensure proper functionality and clean builds with oneAPI.
  • pull/39011

• TPU support for scaled_matmul: The lowering of jax.nn.scaled_matmul is registered for TPU by renaming the ROCm platform-agnostic lowering and adding TPU support, enabling scaled_matmul to run on TPU via an xla.scaled_dot fallback and improving portability without a native TPU fast path.
  • pull/39056

• API documentation categorization: The jax.numpy API documentation is categorized to improve browsing of related functionality, enhancing user experience when exploring the API.
  • pull/39067

• Array API version update and test skips: The array API version is updated to 2025.12, addressing incompatibilities such as functions returning tuples instead of lists by temporarily skipping these cases in tests, with plans for future fixes.
  • pull/39073

• Fixes to jax.numpy functions: Several fixes improve compatibility and behavior of jax.numpy functions: support for negative split indices in jax.numpy.split is restored, jax.numpy.linalg.cond returns infinity instead of NaN for singular matrices, jnp.tri defaults to canonical float dtype respecting jax_enable_x64, and jax.numpy.fill_diagonal gains a wrap=True option matching NumPy behavior.
  • pull/39091, pull/39097, pull/39098, pull/39099

• VJP serialization improvements: A new saveable_args keyword argument is added to jax.vjp allowing selective saving of input arguments as residuals by replacing non-saveable ones with sentinels, enabling more efficient serialization and offloading of VJP objects while requiring explicit restoration before application.
  • pull/39104

• Bug fix in hijax linearize: A bug in the hijax linearize function related to non-differentiable arguments is fixed, improving robustness of differentiation.
  • pull/39127

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

Key Closed Pull Requests

1. Fix static_argnums values leaking into a global jit cache (#16226): This pull request addresses a memory leak in the JAX library by modifying the global LRU cache key in _process_in_axis_resources to exclude static argument values via a new strip_statics() method on FlatTree, thereby preventing user objects passed as static arguments from being indefinitely retained in memory and reducing spurious cache misses without altering value-based caching semantics while functions are alive.

• URL: pull/38947

• Associated Commits: f8d4e, 9883b

• Associated Commits: f8d4e, 9883b

2. Add custom Pallas Xoshiro128++ and Threefry GPU PRNG kernels: This pull request proposes adding custom GPU kernels for the Pallas backend implementing the Xoshiro128++ and Threefry pseudorandom number generators, modeled after existing TPU PRNGs to maintain API consistency and improve performance for stateful random number generation in sequential workloads.

• URL: pull/37190

• Associated Commits: 191ca

• Associated Commits: 191ca

3. [ROCm] Move ROCm Bazel presubmit workflow to 32-bit mode.: This pull request changes the ROCm Bazel presubmit workflow from 64-bit mode to 32-bit mode by setting enable-x64 to 0 in the workflow file to allow for broader testing, with plans to eventually enable both modes for comprehensive testing.

• URL: pull/37475

• Associated Commits: ede81

• Associated Commits: ede81

Other Closed Pull Requests

• Automatic Differentiation Enhancements: Multiple pull requests focus on improving automatic differentiation (AD) in JAX, including the proposal of a separate channel of structured residuals to AD and several fixes in the custom_vjp3 implementation. These changes address tangent handling, error reporting, pretty-printing, and unreachable code removal to enhance AD robustness and usability.
  • pull/37758, pull/39004, pull/39009, pull/39015, pull/39039, pull/39040, pull/39005, pull/39019, pull/39004, pull/39009, pull/39015, pull/39039, pull/39040, pull/39005, pull/39019

• GPU and Hardware Support Improvements: Several pull requests add or improve GPU runtime support, including the addition of OneAPI (SYCL) GPU runtime, enabling mixed FP8 operand types for Mosaic GPU wgmma instructions, and fixing GPU info retrieval on real devices. These updates enhance hardware compatibility and functionality for GPU computations in JAX.
  • pull/38134, pull/38859, pull/38925, pull/38946, pull/38884, pull/38859, pull/38936

• ROCm Build and Testing Enhancements: Multiple pull requests update ROCm-related build configurations and testing workflows, including referencing the latest rules_ml_toolchain, adding a blocking PR gate with curated GPU tests, and replacing hardcoded test exclusions with configurable patterns. These changes improve build reliability and test management for ROCm support.
  • pull/38592, pull/38808, pull/38884

• Memory Management and Allocator Improvements: A pull request enables forwarding the environment variable XLA_PYTHON_CLIENT_ALLOCATOR=address to the PJRT GPU plugin, allowing selection of a dedicated synchronous passthrough allocator. This improves memory management and debugging capabilities for GPU computations in JAX.
  • pull/38490

• Functionality Fixes and API Improvements: Several pull requests fix issues such as incorrect precision handling in jnp.max for bfloat16 inputs, acceptance of negative axis arguments in jax.lax.psum_scatter and jax.lax.all_to_all, and preservation of dictionary key insertion order in pytree operations. These fixes enhance correctness, API consistency, and usability.
  • pull/38984, pull/38998, pull/39020, pull/38950

• Documentation and Educational Content Updates: Pull requests update hijax types documentation by adding mentions of the scan function, extending hitypes documentation with container discipline guidance and new primitives, and adding introductory and intermediate educational materials. These efforts improve documentation clarity and learning resources.
  • pull/38925, pull/39009, pull/39015, pull/39052

• Testing and Platform Compatibility Enhancements: Updates include adding a test case for Thor GPU info, fixing a failing PGLE test by updating expected dump counts, and enabling xla_transform_test.py to run on TPU by skipping certain tests and modifying schedules. These changes improve test coverage and platform compatibility.
  • pull/38936, pull/38946, pull/39041

• Ahead-of-Time Compilation and Backend Initialization: A pull request refactors the register_hlo_module_transformation function to support ahead-of-time (AOT) compilation by obtaining the PJRT_Api* directly from the platform, exposing a platform-to-PJRT_Api* getter to Python, and triggering backend initialization within the function. This eliminates the need for lazy hook functionality and streamlines backend setup.
  • pull/39048

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.