JAX v0.4.13

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).