# JAX v0.4.36 — JAX v0.4.36

- Product: JAX (https://whatsnew.fyi/product/jax)
- Vendor: Google
- Date: 2024-12-05
- Version: v0.4.36
- Original notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.4.36
- Permalink: https://whatsnew.fyi/product/jax/releases/v0.4.36

What's New is an index, not a publisher: every entry below links to the vendor's own release notes, which are the authoritative source. Entries are labelled where they are hand-curated sample data, pre-releases, or drawn from a secondary source such as a developer blog.

Reuse: the summaries, labels and curation here are © What's New. Quote freely with attribution and a link back; wholesale republication of the corpus is not permitted — terms: https://whatsnew.fyi/terms. The vendors' own release notes remain their publishers'.

---

- **changed** — JAX tracing machinery now uses stackless dispatch as a function of context rather than both context and data, removing levels, sublevels, post_process_call, new_base_main, and custom_bind
- **removed** — Removed support for jax.experimental.jax2tf.convert with native_serialization=False or enable_xla=False
- **removed** — Removed xb, xc, and xe symbols from jax.interpreters.xla
- **removed** — Removed deprecated module jax.experimental.export
- **removed** — Removed initial argument to jax.nn.softmax and jax.nn.log_softmax
- **changed** — Calling np.asarray on typed PRNG keys now raises an error instead of returning a scalar object array
- **removed** — Removed jax.export.DisabledSafetyCheck.shape_assertions, jax.export.Exported.lowering_platforms, jax.export.Exported.mlir_module_serialization_version, jax.export.Exported.uses_shape_polymorphism, and lowering_platforms kwarg from jax.export.export
- **removed** — Removed symbolic_scope and symbolic_constraints kwargs from jax.export.symbolic_args_specs
- **changed** — Hashing of tracers now raises a TypeError
- **changed** — JAX build CLI now uses a subcommand structure with build and requirements_update subcommands
- **changed** — jax.scipy.linalg.toeplitz now does implicit batching on multi-dimensional inputs
- **changed** — jax.scipy.special.gamma and jax.scipy.special.gammasgn now return NaN for negative integer inputs
- **removed** — Removed jax.clear_backends
- **removed** — Removed __gpu$xla.gpu.triton custom call from export stability guarantees
- **added** — Added compiler_options argument to jax.jit for passing compilation options to XLA
- **changed** — jax.tree_util.register_dataclass now allows metadata fields to be declared inline via dataclasses.field
- **added** — Added jax.numpy.put_along_axis
- **added** — Added GPU support for jax.lax.linalg.eig and related jax.numpy functions jax.numpy.linalg.eig and jax.numpy.linalg.eigvals
- **added** — Added jax_exec_time_optimization_effort and jax_memory_fitting_effort configuration flags to control compiler effort for execution time and memory usage
- **fixed** — Fixed indexing overflow bug in GPU implementations of LU and QR decomposition for batch sizes close to int32 max
- **deprecated** — Deprecated jax.lib.xla_extension.ArrayImpl and jax.lib.xla_client.ArrayImpl in favor of jax.Array
- **deprecated** — Deprecated jax.lib.xla_extension.XlaRuntimeError in favor of jax.errors.JaxRuntimeError

* Breaking Changes
  * This release lands "stackless", an internal change to JAX's tracing
    machinery. We made trace dispatch purely a function of context rather than a
    function of both context and data. This let us delete a lot of machinery for
    managing data-dependent tracing: levels, sublevels, `post_process_call`,
    `new_base_main`, `custom_bind`, and so on. The change should only affect
    users that use JAX internals.

    If you do use JAX internals then you may need to
    update your code (see
    https://github.com/jax-ml/jax/commit/c36e1f7c1ad4782060cbc8e8c596d85dfb83986f
    for clues about how to do this). There might also be version skew
    issues with JAX libraries that do this. If you find this change breaks your
    non-JAX-internals-using code then try the
    `config.jax_data_dependent_tracing_fallback` flag as a workaround, and if
    you need help updating your code then please file a bug.
  * `jax.experimental.jax2tf.convert` with `native_serialization=False`
    or with `enable_xla=False` have been deprecated since July 2024, with
    JAX version 0.4.31. Now we removed support for these use cases. `jax2tf`
    with native serialization will still be supported.
  * In `jax.interpreters.xla`, the `xb`, `xc`, and `xe` symbols have been removed
    after being deprecated in JAX v0.4.31. Instead use `xb = jax.lib.xla_bridge`,
    `xc = jax.lib.xla_client`, and `xe = jax.lib.xla_extension`.
  * The deprecated module `jax.experimental.export` has been removed. It was replaced
    by `jax.export` in JAX v0.4.30. See the [migration guide](https://jax.readthedocs.io/en/latest/export/export.html#migration-guide-from-jax-experimental-export)
    for information on migrating to the new API.
  * The `initial` argument to `jax.nn.softmax` and `jax.nn.log_softmax`
    has been removed, after being deprecated in v0.4.27.
  * Calling `np.asarray` on typed PRNG keys (i.e. keys produced by `jax.random.key`)
    now raises an error. Previously, this returned a scalar object array.
  * The following deprecated methods and functions in `jax.export` have
    been removed:
      * `jax.export.DisabledSafetyCheck.shape_assertions`: it had no effect
        already.
      * `jax.export.Exported.lowering_platforms`: use `platforms`.
      * `jax.export.Exported.mlir_module_serialization_version`:
        use `calling_convention_version`.
      * `jax.export.Exported.uses_shape_polymorphism`:
         use `uses_global_constants`.
      * the `lowering_platforms` kwarg for `jax.export.export`: use
        `platforms` instead.
  * The kwargs `symbolic_scope` and `symbolic_constraints` from
    `jax.export.symbolic_args_specs` have been removed. They were
    deprecated in June 2024. Use `scope` and `constraints` instead.
  * Hashing of tracers, which has been deprecated since version 0.4.30, now
    results in a `TypeError`.
  * Refactor: JAX build CLI (build/build.py) now uses a subcommand structure and
    replaces previous build.py usage. Run `python build/build.py --help` for
    more details. Brief overview of the new subcommand options:
    * `build`: Builds JAX wheel packages. For e.g., `python build/build.py build --wheels=jaxlib,jax-cuda-pjrt`
    * `requirements_update`: Updates requirements_lock.txt files.
  * `jax.scipy.linalg.toeplitz` now does implicit batching on multi-dimensional
    inputs. To recover the previous behavior, you can call `jax.numpy.ravel`
    on the function inputs.
  * `jax.scipy.special.gamma` and `jax.scipy.special.gammasgn` now
    return NaN for negative integer inputs, to match the behavior of SciPy from
    https://github.com/scipy/scipy/pull/21827.
  * `jax.clear_backends` was removed after being deprecated in v0.4.26.
  * We removed the custom call "__gpu$xla.gpu.triton" from the list of custom
    call that we guarantee export stability. This is because this custom call
    relies on Triton IR, which is not guaranteed t

_[Truncated at 4000 characters — full notes: https://github.com/jax-ml/jax/releases/tag/jax-v0.4.36]_