Skip to content

JAX v0.11.0

Latest
Latest

Choose a tag to compare

View all tags
@hawkinsp hawkinsp released this 16 Jul 19:53
jax-v0.11.0
a152174

  • New features

    • Added a doc on defining custom derivative rules with the experimental
      hijax API (hijax-custom-derivatives), along with
      jax.experimental.hijax helpers for deriving VJPHiPrimitive autodiff
      rules from a jvp or lin rule: linearize_from_jvp with
      apply_derived_linearization, vjp_fwd_from_jvp with transpose_jvp,
      vjp_fwd_from_lin with transpose_linearized, and jvp_from_lin.
    • Added jax.custom_remat to the top-level jax namespace, for
      per-function control of rematerialization under the new jax_remat3
      implementation.
    • jax.checkpoint_policies is now a submodule rather than a namespace
      object (so from jax.checkpoint_policies import ... now works; attribute
      access is unchanged), and it additionally exposes the name-based policy
      classes SaveOnlyTheseNames, SaveAnyNamesButThese, and
      SaveAndOffloadOnlyTheseNames.
    • Added jax.Inline enum for specify inlining policies to
      jax.jit.

  • Breaking changes

    • The deprecated module jax.cloud_tpu_init was removed. This did nothing and
      references to it can be safely removed.
    • Support for Python 3.11, NumPy 2.0, and SciPy 1.14 has been dropped, per the
      deprecation policy.
    • Support for Python 3.13 free-threaded (3.13t) has been dropped. Python
      3.13 free-threaded was an experimental build needed to bootstrap free
      threading support. Now that Python 3.14t is stable, it is time to drop the
      experimental build. Other parts of the Python ecosystem (e.g.
      cibuildwheel, scipy) are making similar moves.
    • jax.numpy.empty and jax.numpy.empty_like now produce
      uninitialized arrays, similar to their NumPy counterparts. Prior to v0.11.0,
      they produced arrays initialized to zeros. To recover the previous behavior,
      use jax.numpy.zeros or jax.numpy.zeros_like instead.

  • Deprecations

    • Passing 2-dimensional arrays (or mixed 2D and 3D arrays) to jax.numpy.cross is deprecated and will be removed in JAX 0.12.0, aligning with NumPy 2.5 behavior.
    • Several previously-deprecated APIs from jax.core have been removed, including
      CallPrimitive, DebugInfo, DropVar, Effect, Effects, InconclusiveDimensionOperation,
      JaxprTypeError, abstract_token, check_jaxpr, concrete_or_error, find_top_trace, gensym,
      get_opaque_trace_state, is_concrete, is_constant_dim, is_constant_shape, jaxprs_in_params,
      new_jaxpr_eqn, no_effects, nonempty_axis_env_DO_NOT_USE, primal_dtype_to_tangent_dtype,
      unsafe_am_i_under_a_jit_DO_NOT_USE, unsafe_am_i_under_a_vmap_DO_NOT_USE,
      unsafe_get_axis_names_DO_NOT_USE, valid_jaxtype, JaxprPpContext, JaxprPpSettings, OutputType,
      aval_mapping_handlers, call, concretization_function_error, custom_typechecks,
      literalable_types, no_axis_name, and trace_ctx.
    • Several previously-deprecated APIs from jax.interpreters.pxla have been removed, including
      Index, MeshAxisName, MeshExecutable, global_aval_to_result_handler, global_result_handlers,
      are_hlo_shardings_equal, is_hlo_sharding_replicated, ArrayMapping, _UNSPECIFIED,
      array_mapping_to_axis_resources, and op_sharding_to_indices.
Assets 2
Loading