JAX v0.11.0

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.