●Change8

Migrating to JAX jax-v0.8.2

Version jax-v0.8.2 introduces 2 breaking changes. This guide details how to update your code.

Released: 12/18/2025

2
Breaking Changes
6
Migration Steps
18
Affected Symbols

⚠️ Check Your Code

If you use any of these symbols, you need to read this guide:

jax.lax.pvaryjax.lax.pcastjax.numpy.arangejax.core.call_impljax.core.get_avaljax.core.mapped_avaljax.core.subjaxprsjax.core.set_current_tracejax.core.take_current_tracejax.core.traverse_jaxpr_paramsjax.core.unmapped_avaljax.core.AbstractTokenjax.core.TraceTagjax.interpreters.pxlajax.Tracerjax.Arrayjax.experimental.si_vjpjax.vjp

Breaking Changes

●Issue #1

jax.Tracer no longer inherits from jax.Array at runtime, which may break code that relies on subclass checks or isinstance behavior for Tracer objects.

●Issue #2

jax.experimental.si_vjp has been removed; code using this function must switch to jax.vjp.

Migration Steps

  1. 1
    Replace any use of `jax.experimental.si_vjp` with `jax.vjp`.
  2. 2
    Update code that imports or calls `jax.lax.pvary` to use `jax.lax.pcast(..., to='varying')`.
  3. 3
    Avoid passing complex numbers to `jax.numpy.arange` or refactor to handle real-valued inputs.
  4. 4
    Remove or replace usage of deprecated symbols from `jax.core` (e.g., `call_impl`, `get_aval`, etc.).
  5. 5
    Stop relying on any APIs from `jax.interpreters.pxla` as they are deprecated.
  6. 6
    If your code checks `issubclass(jax.Tracer, jax.Array)` at runtime, adjust the logic because `Tracer` no longer inherits from `Array`; rely on `isinstance(x, jax.Array)` which still works for traced arrays.

Release Summary

This release deprecates several core and interpreter symbols, removes `jax.experimental.si_vjp`, and changes `Tracer` inheritance, requiring code updates to use `jax.vjp` and the new `pcast` API.

Need More Details?

View the full release notes and all changes for JAX jax-v0.8.2.

View Full Changelog