v0.7.1
JAX v0.7.1
Added 2
- JAX now ships Python 3.14 and 3.14t wheels
- JAX now ships Python 3.13t and 3.14t wheels on Mac
Changed 3
- Exposed jax.set_mesh which acts as a global setter and a context manager
- JAX is now built using CUDA 12.9
- jax.lax.dot now implements the general dot product via the optional dimension_numbers argument
Removed 1
- Removed jax.sharding.use_mesh in favor of jax.set_mesh
Deprecated 4
- jax.lax.zeros_like_array is deprecated in favor of jax.numpy.zeros_like
- Attempting to import jax.experimental.host_callback now results in a DeprecationWarning and will result in an ImportError starting in JAX v0.8.0
- In jax.lax.dot, passing the precision and preferred_element_type arguments by position is deprecated; pass them by explicit keyword instead
- Several internal APIs have been deprecated from jax.interpreters.ad, jax.interpreters.batching, and jax.interpreters.partial_eval
-
New features
- JAX now ships Python 3.14 and 3.14t wheels.
- JAX now ships Python 3.13t and 3.14t wheels on Mac. Previously we only offered free-threading builds on Linux.
-
Changes
- Exposed
jax.set_meshwhich acts as a global setter and a context manager. Removedjax.sharding.use_meshin favor ofjax.set_mesh. - JAX is now built using CUDA 12.9. All versions of CUDA 12.1 or newer remain supported.
jax.lax.dotnow implements the general dot product via the optionaldimension_numbersargument.
- Exposed
-
Deprecations:
jax.lax.zeros_like_arrayis deprecated. Please usejax.numpy.zeros_likeinstead.- Attempting to import
jax.experimental.host_callbacknow results in aDeprecationWarning, and will result in anImportErrorstarting in JAX v0.8.0. Its APIs have raisedNotImplementedErrorsince JAX version 0.4.35. - In
jax.lax.dot, passing theprecisionandpreferred_element_typearguments by position is deprecated. Pass them by explicit keyword instead. - Several dozen internal APIs have been deprecated from
jax.interpreters.ad,jax.interpreters.batching, andjax.interpreters.partial_eval; they are used rarely if ever outside JAX itself, and most are deprecated without any public replacement.