We anticipate that this will be the last release of JAX and jaxlib
supporting a monolithic CUDA jaxlib. Future releases will use the CUDA
plugin jaxlib (e.g. pip install jax[cuda12]).
JAX now requires ml_dtypes version 0.4.0 or newer.
Removed backwards-compatibility support for old usage of the
jax.experimental.export API. It is not possible anymore to use
from jax.experimental.export import export, and instead you should use
from jax.experimental import export.
The removed functionality has been deprecated since 0.4.24.
Deprecations
jax.sharding.XLACompatibleSharding is deprecated. Please use
jax.sharding.Sharding.
jax.experimental.Exported.in_shardings has been renamed as
jax.experimental.Exported.in_shardings_hlo. Same for out_shardings.
The old names will be removed after 3 months.
Removed a number of previously-deprecated APIs:
from {mod}jax.core: non_negative_dim, DimSize, Shape
from {mod}jax.lax: tie_in
from {mod}jax.nn: normalize
from {mod}jax.interpreters.xla: backend_specific_translations,
translations, register_translation, xla_destructure,
TranslationRule, TranslationContext, XlaOp.
The tol argument of {func}jax.numpy.linalg.matrix_rank is being
deprecated and will soon be removed. Use rtol instead.
The rcond argument of {func}jax.numpy.linalg.pinv is being
deprecated and will soon be removed. Use rtol instead.
The deprecated jax.config submodule has been removed. To configure JAX
use import jax and then reference the config object via jax.config.
{mod}jax.random APIs no longer accept batched keys, where previously
some did unintentionally. Going forward, we recommend explicit use of
{func}jax.vmap in such cases.
New Functionality
Added {func}jax.experimental.Exported.in_shardings_jax to construct
shardings that can be used with the JAX APIs from the HloShardings
that are stored in the Exported objects.
Jaxlib release v0.4.29
Bug fixes
Fixed a bug where XLA sharded some concatenation operations incorrectly,
which manifested as an incorrect output for cumulative reductions (#21403).
Fixed a bug where XLA:CPU miscompiled certain matmul fusions
(openxla/xla#13301).
We anticipate that this will be the last release of JAX and jaxlib
supporting a monolithic CUDA jaxlib. Future releases will use the CUDA
plugin jaxlib (e.g. pip install jax[cuda12]).
JAX now requires ml_dtypes version 0.4.0 or newer.
Removed backwards-compatibility support for old usage of the
jax.experimental.export API. It is not possible anymore to use
from jax.experimental.export import export, and instead you should use
from jax.experimental import export.
The removed functionality has been deprecated since 0.4.24.
Added is_leaf argument to {func}jax.tree.all & {func}jax.tree_util.tree_all.
Deprecations
jax.sharding.XLACompatibleSharding is deprecated. Please use
jax.sharding.Sharding.
jax.experimental.Exported.in_shardings has been renamed as
jax.experimental.Exported.in_shardings_hlo. Same for out_shardings.
The old names will be removed after 3 months.
Removed a number of previously-deprecated APIs:
from {mod}jax.core: non_negative_dim, DimSize, Shape
from {mod}jax.lax: tie_in
from {mod}jax.nn: normalize
from {mod}jax.interpreters.xla: backend_specific_translations,
translations, register_translation, xla_destructure,
TranslationRule, TranslationContext, XlaOp.
The tol argument of {func}jax.numpy.linalg.matrix_rank is being
deprecated and will soon be removed. Use rtol instead.
The rcond argument of {func}jax.numpy.linalg.pinv is being
deprecated and will soon be removed. Use rtol instead.
The deprecated jax.config submodule has been removed. To configure JAX
use import jax and then reference the config object via jax.config.
{mod}jax.random APIs no longer accept batched keys, where previously
some did unintentionally. Going forward, we recommend explicit use of
{func}jax.vmap in such cases.
In {func}jax.scipy.special.beta, the x and y parameters have been
renamed to a and b for consistency with other beta APIs.
New Functionality
Added {func}jax.experimental.Exported.in_shardings_jax to construct
shardings that can be used with the JAX APIs from the HloShardings
that are stored in the Exported objects.
jaxlib 0.4.29 (June 10, 2024)
Bug fixes
Fixed a bug where XLA sharded some concatenation operations incorrectly,
which manifested as an incorrect output for cumulative reductions (#21403).
Fixed a bug where XLA:CPU miscompiled certain matmul fusions
... (truncated)
Commits
991797a Merge pull request #21765 from hawkinsp:release
5e7ad60 Removed the double re-exporting of Pallas GPU/TPU APIs
3b4039c [Mosaic GPU] Load LLVM lowering interfaces for all dialects
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.28 to 0.4.29.
Release notes
Sourced from jax's releases.
... (truncated)
Changelog
Sourced from jax's changelog.
... (truncated)
Commits
991797a
Merge pull request #21765 from hawkinsp:release5e7ad60
Removed the double re-exporting of Pallas GPU/TPU APIs3b4039c
[Mosaic GPU] Load LLVM lowering interfaces for all dialectse071053
Prepare for 0.4.29 release.2ade7e7
[pallas] Move the hardware_generation query in the code path that needs itaf95803
Merge pull request #21759 from rajasekharporeddy:testbranch1775c6f8
Fix Typos in docs and one error message8fbe65b
Update XLA dependency to use revision6617a0d
Expanddevice_put
benchmarks to run with different numbers of arrays and in...a8246ea
Issue a warning where code relies on a bug where treedef.flatten_up_to(...) w...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