This release includes wheels for Python 3.13. Free-threading mode is not yet
supported.
jax.errors.JaxRuntimeError has been added as a public alias for the
formerly private XlaRuntimeError type.
Breaking changes
jax_pmap_no_rank_reduction flag is set to True by default.
array[0] on a pmap result now introduces a reshape (use array[0:1]
instead).
The per-shard shape (accessable via jax_array.addressable_shards or
jax_array.addressable_data(0)) now has a leading (1, ...). Update code
that directly accesses shards accordingly. The rank of the per-shard-shape
now matches that of the global shape which is the same behavior as jit.
This avoids costly reshapes when passing results from pmap into jit.
jax.experimental.host_callback has been deprecated since March 2024, with
JAX version 0.4.26. Now we set the default value of the
--jax_host_callback_legacy configuration value to True, which means that
if your code uses jax.experimental.host_callback APIs, those API calls
will be implemented in terms of the new jax.experimental.io_callback API.
If this breaks your code, for a very limited time, you can set the
--jax_host_callback_legacy to True. Soon we will remove that
configuration option, so you should instead transition to using the
new JAX callback APIs. See #20385 for a discussion.
Deprecations
In jax.numpy.trim_zeros, non-arraylike arguments or arraylike
arguments with ndim != 1 are now deprecated, and in the future will result
in an error.
Internal pretty-printing tools jax.core.pp_* have been removed, after
being deprecated in JAX v0.4.30.
jax.lib.xla_client.Device is deprecated; use jax.Device instead.
jax.lib.xla_client.XlaRuntimeError has been deprecated. Use
jax.errors.JaxRuntimeError instead.
The default behavior of jax.pure_callback and
jax.extend.ffi.ffi_call under vmap has been deprecated and so has
the vectorized parameter to those functions. The vmap_method parameter
should be used instead for better defined behavior. See the discussion in
#23881 for more details.
Deletion:
jax.xla_computation is deleted. It has been 3 months since its deprecation
in 0.4.30 JAX release.
Please use the AOT APIs to get the same functionality as jax.xla_computation.
jax.xla_computation(fn)(*args, **kwargs) can be replaced with
jax.jit(fn).lower(*args, **kwargs).compiler_ir('hlo').
You can also use .out_info property of jax.stages.Lowered to get the
output information (like tree structure, shape and dtype).
This release includes wheels for Python 3.13. Free-threading mode is not yet
supported.
jax.errors.JaxRuntimeError has been added as a public alias for the
formerly private XlaRuntimeError type.
Breaking changes
jax_pmap_no_rank_reduction flag is set to True by default.
array[0] on a pmap result now introduces a reshape (use array[0:1]
instead).
The per-shard shape (accessable via jax_array.addressable_shards or
jax_array.addressable_data(0)) now has a leading (1, ...). Update code
that directly accesses shards accordingly. The rank of the per-shard-shape
now matches that of the global shape which is the same behavior as jit.
This avoids costly reshapes when passing results from pmap into jit.
jax.experimental.host_callback has been deprecated since March 2024, with
JAX version 0.4.26. Now we set the default value of the
--jax_host_callback_legacy configuration value to True, which means that
if your code uses jax.experimental.host_callback APIs, those API calls
will be implemented in terms of the new jax.experimental.io_callback API.
If this breaks your code, for a very limited time, you can set the
--jax_host_callback_legacy to True. Soon we will remove that
configuration option, so you should instead transition to using the
new JAX callback APIs. See {jax-issue}[#20385](https://github.com/jax-ml/jax/issues/20385) for a discussion.
Deprecations
In {func}jax.numpy.trim_zeros, non-arraylike arguments or arraylike
arguments with ndim != 1 are now deprecated, and in the future will result
in an error.
Internal pretty-printing tools jax.core.pp_* have been removed, after
being deprecated in JAX v0.4.30.
jax.lib.xla_client.Device is deprecated; use jax.Device instead.
jax.lib.xla_client.XlaRuntimeError has been deprecated. Use
jax.errors.JaxRuntimeError instead.
The default behavior of {func}jax.pure_callback and
{func}jax.extend.ffi.ffi_call under vmap has been deprecated and so has
the vectorized parameter to those functions. The vmap_method parameter
should be used instead for better defined behavior. See the discussion in
{jax-issue}[#23881](https://github.com/jax-ml/jax/issues/23881) for more details.
Deletion:
jax.xla_computation is deleted. It's been 3 months since it's deprecation
in 0.4.30 JAX release.
Please use the AOT APIs to get the same functionality as jax.xla_computation.
jax.xla_computation(fn)(*args, **kwargs) can be replaced with
jax.jit(fn).lower(*args, **kwargs).compiler_ir('hlo').
You can also use .out_info property of jax.stages.Lowered to get the
output information (like tree structure, shape and dtype).
Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting @dependabot rebase.
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR:
- `@dependabot rebase` will rebase this PR
- `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it
- `@dependabot merge` will merge this PR after your CI passes on it
- `@dependabot squash and merge` will squash and merge this PR after your CI passes on it
- `@dependabot cancel merge` will cancel a previously requested merge and block automerging
- `@dependabot reopen` will reopen this PR if it is closed
- `@dependabot close` will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually
- `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency
- `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
- `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
- `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
Bumps jax from 0.4.33 to 0.4.34.
Release notes
Sourced from jax's releases.
... (truncated)
Changelog
Sourced from jax's changelog.
... (truncated)
Commits
affba36
JAX 0.4.34 release59cb692
Update XLA dependency to use revision8fc8082
Merge pull request #24077 from jakevdp:badge-linkscfb7541
Merge pull request #23738 from keshavb96:disable_remat_pass78b65dd
Merge pull request #24070 from jakevdp:issubdtype-doca8744e3
README: fix badge href links152a873
Update XLA dependency to use revisionc9f946e
Only thread a discharged ref value through a cond when it changes in some branch816947b
Merge pull request #24071 from hawkinsp:wincie212c77
Merge pull request #23891 from ROCm:build-fixes-rollupDependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting
@dependabot rebase
.Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot merge` will merge this PR after your CI passes on it - `@dependabot squash and merge` will squash and merge this PR after your CI passes on it - `@dependabot cancel merge` will cancel a previously requested merge and block automerging - `@dependabot reopen` will reopen this PR if it is closed - `@dependabot close` will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually - `@dependabot show