Releases · google/jax

NOTE: This is the last JAX release that will include Python 3.8 support

Changes
- jax.jit now allows None to be passed to in_shardings and
  out_shardings. The semantics are as follows:
  - For in_shardings, JAX will mark is as replicated but this behavior
    can change in the future.
  - For out_shardings, we will rely on the XLA GSPMD partitioner to
    determine the output shardings.
- jax.experimental.pjit.pjit also allows None to be passed to
  in_shardings and out_shardings. The semantics are as follows:
  - If the mesh context manager is not provided, JAX has the freedom to
    choose whatever sharding it wants.
    - For in_shardings, JAX will mark is as replicated but this behavior
      can change in the future.
    - For out_shardings, we will rely on the XLA GSPMD partitioner to
      determine the output shardings.
  - If the mesh context manager is provided, None will imply that the value
    will be replicated on all devices of the mesh.
- Executable.cost_analysis() works on Cloud TPU
- Added a warning if a non-allowlisted jaxlib plugin is in use.
- Added jax.tree_util.tree_leaves_with_path.
Bug fixes
- Fixed incorrect wheel name in CUDA 12 releases (#16362); the correct wheel
  is named cudnn89 instead of cudnn88.
Deprecations
- The native_serialization_strict_checks parameter to
  {func}jax.experimental.jax2tf.convert is deprecated in favor of the
  new native_serializaation_disabled_checks ({jax-issue}#16347).

Provide feedback