Data Science, Python & R — July 17, 2026

Executive summary

JAX v0.11.0 (2026-07-16) introduces advanced developer-facing autodiff features (custom-derivative helpers, top-level custom_remat, checkpoint policy submodule, inlining control) while making multiple breaking and compatibility changes: removal of older Python/NumPy/SciPy support, behavioral change for jax.numpy.empty/empty_like, deprecation of certain jax.numpy.cross usages, and wholesale removal of many previously-deprecated internals — signaling a push toward a smaller, more maintainable public surface and closer alignment with evolving NumPy behavior [1].

Chronological timeline of key developments

• 2026-07-16 — JAX v0.11.0 released:
  • New APIs & docs for custom derivative rules via experimental jax.experimental.hijax and helper primitives (linearize_from_jvp, apply_derived_linearization, vjp_fwd_from_jvp, transpose_jvp, etc.) [1].
  • Added jax.custom_remat at top-level and a checkpoint_policies submodule exposing name-based policies (SaveOnlyTheseNames, SaveAnyNamesButThese, SaveAndOffloadOnlyTheseNames) [1].
  • Added jax.Inline enum to control inlining behavior in jax.jit [1].
  • Breaking and compatibility changes: removed jax.cloud_tpu_init; dropped support for Python 3.11, NumPy 2.0, SciPy 1.14; changed free-threaded build policy now that Python 3.14t is stable; jax.numpy.empty/empty_like now return uninitialized arrays (previously zeros) [1].
  • Deprecation: passing 2D (or mixed 2D/3D) arrays to jax.numpy.cross is deprecated and will be removed in JAX 0.12.0 to match NumPy 2.5 behavior [1].
  • Removed many deprecated internals from jax.core and jax.interpreters.pxla that downstream code may have relied on [1].

Trends

• API maturation and surface shrinking: Removal of long-deprecated internals indicates a stabilization effort and a cleaner public API surface, reducing maintenance burden but increasing migration needs for downstream projects [1].
• Tighter alignment with NumPy evolution: Deprecations (e.g., jax.numpy.cross) and compatibility notes reflect an effort to track NumPy 2.x/2.5 behavior, implying future JAX releases will continue to follow NumPy semantics closely [1].
• Performance and developer control: New features (custom derivative helpers, custom_remat, checkpoint policies, inlining control) emphasize enabling finer-grained control for performance and advanced differentiation workflows [1].
• Platform/support consolidation: Dropping older Python/NumPy/SciPy versions indicates consolidation around more recent language and library releases to simplify CI and optimize builds [1].

Risks

• Compatibility breaks for user code and libraries: Removal of deprecated internals and dropped platform versions will break code that relied on internal APIs or the older supported runtimes; downstream ecosystems (libraries, frameworks) may need rapid updates [1].
• Behavioral changes causing subtle bugs: jax.numpy.empty/empty_like returning uninitialized memory (instead of zeros) can introduce silent correctness bugs if users assume zero-initialization [1].
• Migration burden: Deprecation of certain argument shapes (jax.numpy.cross) and upcoming removals in 0.12.0 force code changes across many projects, increasing short-term workload [1].
• CI and packaging complexity: Dropped support for Python 3.11/NumPy 2.0/SciPy 1.14 may require updates to pinned dependency matrices, container images, and platform support statements [1].

Opportunities

• Adopt new derivative and rematerialization APIs: Use the hijax helpers and jax.custom_remat to implement more efficient/custom autodiff logic and reduce memory/compute trade-offs in models [1].
• Fine-tune performance: Leverage jax.Inline and the checkpoint_policies submodule to tune JIT inlining and checkpoint behavior for critical workloads.
• Improve code quality: Removing dependency on deprecated internals is an opportunity to migrate to stable public APIs, improving long-term maintainability.
• Align with NumPy-driven ecosystem: Early adaptation to NumPy semantics changes reduces future churn and positions projects to benefit from NumPy performance/semantics improvements [1].

Recommended actions

• Immediate compatibility audit: Scan codebases and CI for uses of removed internals (jax.core and jax.interpreters.pxla symbols) and replace with public APIs or upstream-compatible alternatives; run test suites under JAX v0.11.0 to identify breakage [1].
• Fix code assuming zero-initialized arrays: Replace uses of jax.numpy.empty/empty_like where zero-initialization was expected with jax.numpy.zeros/zeros_like, and audit any buffer-initialization logic for correctness [1].
• Address deprecated jax.numpy.cross usage: Find code paths passing 2D (or mixed 2D/3D) arrays to jax.numpy.cross and refactor to supported shapes ahead of JAX 0.12.0 [1].
• Update dependency and platform matrices: Adjust pinned Python/NumPy/SciPy versions in manifests, Docker images, and CI to reflect dropped support (remove Python 3.11/NumPy 2.0/SciPy 1.14 from supported lists) or pin to earlier JAX releases if immediate upgrade is infeasible [1].
• Leverage new features where beneficial: Identify hotspots (memory/compute trade-offs, complex custom ops) that can benefit from hijax helpers, jax.custom_remat, checkpoint policies, and jax.Inline; prototype and measure gains in staging environments [1].
• Communicate with downstream maintainers: If you maintain libraries that target JAX, announce required changes, provide migration examples, and consider temporary compatibility layers to assist users.
• Monitor upcoming releases: Track JAX 0.12.0 for final removals and NumPy 2.5 compatibility changes to plan subsequent migrations [1].

Sources

1. [1] JAX v0.11.0