Support for Python 3.7 has been dropped, in accordance with JAX's
{ref}version-support-policy.
We introduce jax.Array which is a unified array type that subsumes
DeviceArray, ShardedDeviceArray, and GlobalDeviceArray types in JAX.
The jax.Array type helps make parallelism a core feature of JAX,
simplifies and unifies JAX internals, and allows us to unify jit and
pjit. jax.Array has been enabled by default in JAX 0.4 and makes some
breaking change to the pjit API. The jax.Array migration
guide can
help you migrate your codebase to jax.Array. You can also look at the
Distributed arrays and automatic parallelization
tutorial to understand the new concepts.
PartitionSpec and Mesh are now out of experimental. The new API endpoints
are jax.sharding.PartitionSpec and jax.sharding.Mesh.
jax.experimental.maps.Mesh and jax.experimental.PartitionSpec are
deprecated and will be removed in 3 months.
with_sharding_constraints new public endpoint is
jax.lax.with_sharding_constraint.
If using ABSL flags together with jax.config, the ABSL flag values are no
longer read or written after the JAX configuration options are initially
populated from the ABSL flags. This change improves performance of reading
jax.config options, which are used pervasively in JAX.
The jax2tf.call_tf function now uses for TF lowering the first TF
device of the same platform as used by the embedding JAX computation.
Before, it was using the 0th device for the JAX-default backend.
A number of jax.numpy functions now have their arguments marked as
positional-only, matching NumPy.
jnp.msort is now deprecated, following the deprecation of np.msort in numpy 1.24.
It will be removed in a future release, in accordance with the {ref}api-compatibility
policy. It can be replaced with jnp.sort(a, axis=0).
Jaxlib release v0.4.1
Changes
Support for Python 3.7 has been dropped, in accordance with JAX's
{ref}version-support-policy.
The behavior of XLA_PYTHON_CLIENT_MEM_FRACTION=.XX has been changed to allocate XX% of
Deleted {func}jax.scipy.linalg.polar_unitary, which was a deprecated JAX
extension to the scipy API. Use {func}jax.scipy.linalg.polar instead.
Changes
Added {func}jax.scipy.stats.rankdata.
jaxlib 0.4.3 (Feb 8, 2023)
jax 0.4.2 (Jan 24, 2023)
Breaking changes
Deleted jax.experimental.callback
Operations with dimensions in presence of jax2tf shape polymorphism have
been generalized to work in more scenarios, by converting the symbolic
dimension to JAX arrays. Operations involving symbolic dimensions and
np.ndarray now can raise errors when the result is used as a shape value
({jax-issue}[#14106](https://github.com/google/jax/issues/14106)).
jaxpr objects now raise an error on attribute setting in order to avoid
problematic mutations ({jax-issue}14102)
Changes
{func}jax2tf.call_tf has a new parameter has_side_effects (default True)
that can be used to declare whether an instance can be removed or replicated
by JAX optimizations such as dead-code elimination ({jax-issue}[#13980](https://github.com/google/jax/issues/13980)).
Added more support for floordiv and mod for jax2tf shape polymorphism. Previously,
certain division operations resulted in errors in presence of symbolic dimensions
({jax-issue}[#14108](https://github.com/google/jax/issues/14108)).
jaxlib 0.4.2 (Jan 24, 2023)
Changes
Set JAX_USE_PJRT_C_API_ON_TPU=1 to enable new Cloud TPU runtime, featuring
automatic device memory defragmentation.
jax 0.4.1 (Dec 13, 2022)
Changes
Support for Python 3.7 has been dropped, in accordance with JAX's
{ref}version-support-policy.
We introduce jax.Array which is a unified array type that subsumes
DeviceArray, ShardedDeviceArray, and GlobalDeviceArray types in JAX.
The jax.Array type helps make parallelism a core feature of JAX,
simplifies and unifies JAX internals, and allows us to unify jit and
pjit. jax.Array has been enabled by default in JAX 0.4 and makes some
breaking change to the pjit API. The jax.Array migration
guide can
help you migrate your codebase to jax.Array. You can also look at the
Distributed arrays and automatic parallelization
... (truncated)
Commits
7783a5e Merge pull request #14343 from skye:cache_options_check
1228cbd Change the executable_build_options check in compilation_cache.py to be more ...
You can trigger a rebase of this PR 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 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.3.25 to 0.4.3.
Release notes
Sourced from jax's releases.
... (truncated)
Changelog
Sourced from jax's changelog.
... (truncated)
Commits
7783a5e
Merge pull request #14343 from skye:cache_options_check1228cbd
Change the executable_build_options check in compilation_cache.py to be more ...a46f31b
Merge pull request #14342 from skye:version8ab1585
Update WORKSPACE and setup.py for jax/jaxlib 0.4.3 releaseeb13c05
Add option to run tests with persistent compilation cache enabled.6860cb8
Move jax.interpreters.xla to jax._src.interpreters.xla.9c827fb
Merge pull request #14340 from ROCmSoftwarePlatform:rocm_reenable_linalg_spar...01a10a1
[ROCm] Re-enable some linalg and sparse tests8eb00c5
Merge pull request #14335 from jakevdp:doc-transformationsa022a4e
DOC: remove transformations.mdYou can trigger a rebase of this PR 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 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)