Bump jax from 0.8.1 to 0.8.2

[dependabot]

Bumps jax from 0.8.1 to 0.8.2.

Release notes

Sourced from jax's releases.

JAX v0.8.2

• Deprecations

  • jax.lax.pvary has been deprecated. Please use jax.lax.pcast(..., to='varying') as the replacement.
  • Complex arguments passed to jax.numpy.arange now result in a deprecation warning, because the output is poorly-defined.
  • From jax.core a number of symbols are newly deprecated including: call_impl, get_aval, mapped_aval, subjaxprs, set_current_trace, take_current_trace, traverse_jaxpr_params, unmapped_aval, AbstractToken, and TraceTag.
  • All symbols in jax.interpreters.pxla are deprecated. These are primarily JAX internal APIs, and users should not rely on them.

• Changes:

  • jax's Tracer no longer inherits from jax.Array at runtime. However, jax.Array now uses a custom metaclass such isinstance(x, Array) is true if an object x represents a traced Array. Only some Tracers represent Arrays, so it is not correct for Tracer to inherit from Array.

    For the moment, during Python type checking, we continue to declare Tracer as a subclass of Array, however we expect to remove this in a future release.

  • jax.experimental.si_vjp has been deleted. jax.vjp subsumes it's functionality.

Changelog

Sourced from jax's changelog.

JAX 0.8.2 (December 18, 2025)

• Deprecations

  • jax.lax.pvary has been deprecated. Please use jax.lax.pcast(..., to='varying') as the replacement.
  • Complex arguments passed to {func}jax.numpy.arange now result in a deprecation warning, because the output is poorly-defined.
  • From {mod}jax.core a number of symbols are newly deprecated including: call_impl, get_aval, mapped_aval, subjaxprs, set_current_trace, take_current_trace, traverse_jaxpr_params, unmapped_aval, AbstractToken, and TraceTag.
  • All symbols in {mod}jax.interpreters.pxla are deprecated. These are primarily JAX internal APIs, and users should not rely on them.

• Changes:

  • jax's Tracer no longer inherits from jax.Array at runtime. However, jax.Array now uses a custom metaclass such isinstance(x, Array) is true if an object x represents a traced Array. Only some Tracers represent Arrays, so it is not correct for Tracer to inherit from Array.

    For the moment, during Python type checking, we continue to declare Tracer as a subclass of Array, however we expect to remove this in a future release.

  • jax.experimental.si_vjp has been deleted. jax.vjp subsumes it's functionality.

Commits

• 3f12502 Prepare for JAX release 0.8.2
• e2815d5 [pallas:mosaic] pltpu.emit_pipeline now accepts block specs in HBM
• 69a6cca [mpmd] Fix stage_id field nanobind compatability
• 315bb93 [MGPU] Doc fix and shape size fix for broadcast in WGSplatFragLayout.
• c1188e7 Add getitem to backwards compatible shims.
• b9b2d4f [pallas:mosaic] Recover memory_space from the aval in aval_to_ir_type
• 3349e7a [Pallas] Allow inferring the backend from the provided compiler params.
• 97c557d Update XLA dependency to use revision http://github.com/openxla/xla/commit/91...
• 286e358 Automated Code Change
• f81c201 [pallas:mosaic] Access memory spaces via pltpu directly
• Additional commits viewable in compare view

You can trigger a rebase of this PR by commenting @dependabot rebase.

---

Dependabot commands and options

You can trigger Dependabot actions by commenting on this PR:

• @dependabot rebase will rebase this PR
• @dependabot recreate will recreate this PR, overwriting any edits that have been made to it
• @dependabot merge will merge this PR after your CI passes on it
• @dependabot squash and merge will squash and merge this PR after your CI passes on it
• @dependabot cancel merge will cancel a previously requested merge and block automerging
• @dependabot reopen will reopen this PR if it is closed
• @dependabot close will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually
• @dependabot show <dependency name> ignore conditions will show all of the ignore conditions of the specified dependency
• @dependabot ignore this major version will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this minor version will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
• @dependabot ignore this dependency will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)

Note
Automatic rebases have been disabled on this pull request as it has been open for over 30 days.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 96.98%. Comparing base (02446c5) to head (9847bbc).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #639   +/-   ##
=======================================
  Coverage   96.98%   96.98%           
=======================================
  Files          42       42           
  Lines        1094     1094           
=======================================
  Hits         1061     1061           
  Misses         33       33           

Flag	Coverage Δ
unittests	96.98% <ø> (ø)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

Thank you for your contribution @dependabot[bot] 🚀! Your github-pages is ready for download 👉 here 👈!
_{(The artifact expires on 2025-12-29T05:23:24Z. You can re-generate it by re-running the workflow here.)}

[dependabot]

OK, I won't notify you again about this release, but will get in touch when a new version is available. If you'd rather skip all updates until the next major or minor version, let me know by commenting @dependabot ignore this major version or @dependabot ignore this minor version. You can also ignore all major, minor, or patch releases for a dependency by adding an ignore condition with the desired update_types to your config file.

If you change your mind, just re-open this PR and I'll resolve any conflicts on it.