mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-14 19:06:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/CHANGELOG.md

3176 lines
157 KiB
Markdown
Raw Permalink Normal View History

move changelog to top level
2021-03-05 11:07:50 -08:00
# Change log
Best viewed [here](https://jax.readthedocs.io/en/latest/changelog.html).
[pallas] Added a CHANGELOG for Pallas The CHANGELOG is populated with the changes since June 10th, when JAX 0.4.29 was released.
2024-07-11 11:55:31 +03:00
For the changes specific to the experimental Pallas APIs,
see {ref}`pallas-changelog`.
move changelog to top level
2021-03-05 11:07:50 -08:00
changelog: link to api compatibility & python version docs
2024-12-23 09:39:45 -08:00
JAX follows Effort-based versioning; for a discussion of this and JAX's API
compatibility policy, refer to {ref}`api-compatibility`. For the Python and
NumPy version support policy, refer to {ref}`version-support-policy`.
move changelog to top level
2021-03-05 11:07:50 -08:00
<!--
Remember to align the itemized text with the first line of an item within a list.
[pallas] Improve the error localization * Add the source location information for the index map function to `BlockMapping`. * Removed the `compute_index` wrapper around the index_map, so that we can get the location information for the index_map, not the wrapper. * Added source location to the errors related to index map functions. * Added an error if the index map returns something other than integer scalars. * Construct BlockSpec origins for arguments using JAX helper functions to get argument names * Removed redundant API error tests from tpu_pallas_test.py
2024-07-18 15:33:40 +02:00
When releasing, please add the new-release-boilerplate to docs/pallas/CHANGELOG.md.
move changelog to top level
2021-03-05 11:07:50 -08:00
-->
Update version numbers after 0.5.1 release.
2025-02-24 16:07:26 -05:00
## Unreleased
Add an allow_negative_indices option to lax.dynamic_slice and lax.dynamic_update_slice. The goal of this change is to avoid generating code to wrap negative indices back into range in cases where we know it doesn't matter. Change scan to pass allow_negative_indices=False to avoid emitting index wrapping code for each scan argument. PiperOrigin-RevId: 731812827
2025-02-27 12:02:13 -08:00
* New Features
* Added a `allow_negative_indices` option to {func}`jax.lax.dynamic_slice`,
{func}`jax.lax.dynamic_update_slice` and related functions. The default is
true, matching the current behavior. If set to false, JAX does not need to
emit code clamping negative indices, which improves code size.
Add replace option to random.categorical to enable sampling without replacement.
2025-03-13 19:00:19 -04:00
* Added a `replace` option to {func}`jax.random.categorical` to enable sampling
without replacement.
Add an allow_negative_indices option to lax.dynamic_slice and lax.dynamic_update_slice. The goal of this change is to avoid generating code to wrap negative indices back into range in cases where we know it doesn't matter. Change scan to pass allow_negative_indices=False to avoid emitting index wrapping code for each scan argument. PiperOrigin-RevId: 731812827
2025-02-27 12:02:13 -08:00
Prepare for JAX release 0.5.2 (patch release over 0.5.1)
2025-03-04 12:50:43 -08:00
## jax 0.5.2 (Mar 4, 2025)
Patch release of 0.5.1
* Bug fixes
* Fixes TPU metric logging and `tpu-info`, which was broken in 0.5.1
Prepare for JAX release 0.5.1
2025-02-18 09:29:24 -08:00
## jax 0.5.1 (Feb 24, 2025)
Update version numbers after 0.5.0 release
2025-01-17 13:30:59 -05:00
Add custom_dce to changelogs and API docs.
2025-01-27 11:16:00 -05:00
* New Features
* Added an experimental {func}`jax.experimental.custom_dce.custom_dce`
decorator to support customizing the behavior of opaque functions under
JAX-level dead code elimination (DCE). See {jax-issue}`#25956` for more
details.
Add public APIs for jax.lax monoidal reductions
2025-02-11 16:00:03 -08:00
* Added low-level reduction APIs in {mod}`jax.lax`: {func}`jax.lax.reduce_sum`,
Don't write atime file if JAX_COMPILATIION_CACHE_MAX_SIZE == -1 The atime file is only needed to implement the LRU eviction policy, which is only needed if a max persistence compilation cache size is set. Writing this file can cause network filesystem performace and other issues, so only write it if users are opted-in.
2025-02-13 17:37:27 -08:00
{func}`jax.lax.reduce_prod`, {func}`jax.lax.reduce_max`, {func}`jax.lax.reduce_min`,
Add public APIs for jax.lax monoidal reductions
2025-02-11 16:00:03 -08:00
{func}`jax.lax.reduce_and`, {func}`jax.lax.reduce_or`, and {func}`jax.lax.reduce_xor`.
Enable pivoted QR on GPU via MAGMA. Originally noted in #20282, this commit provides a GPU compatible implementation of `geqp3` via MAGMA.
2025-01-10 15:55:53 +00:00
* {func}`jax.lax.linalg.qr`, and {func}`jax.scipy.linalg.qr`, now support
column-pivoting on CPU and GPU. See {jax-issue}`#20282` and
Add jax.random.multinomial.
2025-03-12 18:15:14 -04:00
* Added {func}`jax.random.multinomial`.
Enable pivoted QR on GPU via MAGMA. Originally noted in #20282, this commit provides a GPU compatible implementation of `geqp3` via MAGMA.
2025-01-10 15:55:53 +00:00
{jax-issue}`#25955` for more details.
Add custom_dce to changelogs and API docs.
2025-01-27 11:16:00 -05:00
Make `JAX_CPU_COLLECTIVES_IMPLEMENTATION` and `JAX_NUM_CPU_DEVICES` env vars Before, these values could only be specified via jax.config or flags. This PR makes them proper configs, so they also work as env vars.
2025-01-28 17:00:06 -08:00
* Changes
* `JAX_CPU_COLLECTIVES_IMPLEMENTATION` and `JAX_NUM_CPU_DEVICES` now work as
env vars. Before they could only be specified via jax.config or flags.
Default JAX_CPU_COLLECTIVES_IMPLEMENTATION to 'gloo'. This enables CPU collectives by default, making multi-process CPU communication work without extra configuration. PiperOrigin-RevId: 724076284
2025-02-06 14:29:52 -08:00
* `JAX_CPU_COLLECTIVES_IMPLEMENTATION` now defaults to `'gloo'`, meaning
multi-process CPU communication works out-of-the-box.
Remove libtpu-nightly dependency from jax[tpu]. For several releases, libtpu-nightly has been a transitional empty package that does nothing. We remove the dependency in preparation for depending on libtpu from pypi instead of a GCS bucket in jax[tpu].
2025-02-04 20:58:50 -05:00
* The `jax[tpu]` TPU extra no longer depends on the `libtpu-nightly` package.
This package may safely be removed if it is present on your machine; JAX now
uses `libtpu` instead.
Make `JAX_CPU_COLLECTIVES_IMPLEMENTATION` and `JAX_NUM_CPU_DEVICES` env vars Before, these values could only be specified via jax.config or flags. This PR makes them proper configs, so they also work as env vars.
2025-01-28 17:00:06 -08:00
[better_errors] Make it explicit that debug_info is not None. Now all internal uses of lu.wrap_init and core.Jaxpr are with actual debug info. This enables us to clean up the type declarations and to remove the checks whether debug_info is present. For usage outside of the JAX internals, we change `jax.extend.linear_util.wrap_init` to be usable without debug_info, for temporary backwards compatibility. We emit a deprecation warning and fill-in some fake debugging info. See https://github.com/jax-ml/jax/issues/26480 for more details. PiperOrigin-RevId: 726770483
2025-02-13 22:06:18 -08:00
* Deprecations
* The internal function `linear_util.wrap_init` and the constructor
`core.Jaxpr` now must take a non-empty `core.DebugInfo` kwarg. For
a limited time, a `DeprecationWarning` is printed if
`jax.extend.linear_util.wrap_init` is used without debugging info.
A downstream effect of this several other internal functions need debug
info. This change does not affect public APIs.
See https://github.com/jax-ml/jax/issues/26480 for more detail.
jax.numpy ndim/shape/size: deprecate non-array input
2025-02-20 12:31:31 -08:00
* In {func}`jax.numpy.ndim`, {func}`jax.numpy.shape`, and {func}`jax.numpy.size`,
non-arraylike inputs (such as lists, tuples, etc.) are now deprecated.
[better_errors] Make it explicit that debug_info is not None. Now all internal uses of lu.wrap_init and core.Jaxpr are with actual debug info. This enables us to clean up the type declarations and to remove the checks whether debug_info is present. For usage outside of the JAX internals, we change `jax.extend.linear_util.wrap_init` to be usable without debug_info, for temporary backwards compatibility. We emit a deprecation warning and fill-in some fake debugging info. See https://github.com/jax-ml/jax/issues/26480 for more details. PiperOrigin-RevId: 726770483
2025-02-13 22:06:18 -08:00
Don't write atime file if JAX_COMPILATIION_CACHE_MAX_SIZE == -1 The atime file is only needed to implement the LRU eviction policy, which is only needed if a max persistence compilation cache size is set. Writing this file can cause network filesystem performace and other issues, so only write it if users are opted-in.
2025-02-13 17:37:27 -08:00
* Bug fixes
Prepare for JAX release 0.5.1
2025-02-18 09:29:24 -08:00
* TPU runtime startup and shutdown time should be significantly improved on
TPU v5e and newer (from around 17s to around 8s). If not already set, you may
need to enable transparent hugepages in your VM image
(`sudo sh -c 'echo always > /sys/kernel/mm/transparent_hugepage/enabled'`).
We hope to improve this further in future releases.
Don't write atime file if JAX_COMPILATIION_CACHE_MAX_SIZE == -1 The atime file is only needed to implement the LRU eviction policy, which is only needed if a max persistence compilation cache size is set. Writing this file can cause network filesystem performace and other issues, so only write it if users are opted-in.
2025-02-13 17:37:27 -08:00
* Persistent compilation cache no longer writes access time file if
JAX_COMPILATION_CACHE_MAX_SIZE is unset or set to -1, i.e. if the LRU
eviction policy isn't enabled. This should improve performance when using
the cache with large-scale network storage.
Release JAX 0.5.0
2025-01-17 10:28:03 -05:00
## jax 0.5.0 (Jan 17, 2025)
Update version numbers after release.
2024-12-17 18:03:28 -05:00
Update the JAX version to 0.5.0. This is because of the breaking change to PRNG key semantics, and the version follows JAX's new effver versioning scheme (https://jax.readthedocs.io/en/latest/jep/25516-effver.html).
2025-01-15 14:08:15 -05:00
As of this release, JAX now uses
[effort-based versioning](https://jax.readthedocs.io/en/latest/jep/25516-effver.html).
Release JAX 0.5.0
2025-01-17 10:28:03 -05:00
Since this release makes a breaking change to PRNG key semantics that
Update the JAX version to 0.5.0. This is because of the breaking change to PRNG key semantics, and the version follows JAX's new effver versioning scheme (https://jax.readthedocs.io/en/latest/jep/25516-effver.html).
2025-01-15 14:08:15 -05:00
may require users to update their code, we are bumping the "meso" version of JAX
to signify this.
enable partitionable threefry by default PiperOrigin-RevId: 715242560
2025-01-13 22:45:41 -08:00
* Breaking changes
* Enable `jax_threefry_partitionable` by default (see
[the update note](https://github.com/jax-ml/jax/discussions/18480)).
Release JAX 0.5.0
2025-01-17 10:28:03 -05:00
* This release drops support for Mac x86 wheels. Mac ARM of course remains
supported. For a recent discussion, see
https://github.com/jax-ml/jax/discussions/22936.
Two key factors motivated this decision:
* The Mac x86 build (only) has a number of test failures and crashes. We
would prefer to ship no release than a broken release.
* Mac x86 hardware is end-of-life and cannot be easily obtained for
developers at this point. So it is difficult for us to fix this kind of
problem even if we wanted to.
We are open to readding support for Mac x86 if the community is willing
to help support that platform: in particular, we would need the JAX test
suite to pass cleanly on Mac x86 before we could ship releases again.
Increase the minimum NumPy version to v1.25. Per SPEC 0, we drop NumPy v1.24 support on Dec 18, 2024.
2024-12-18 08:18:57 -05:00
* Changes:
* The minimum NumPy version is now 1.25. NumPy 1.25 will remain the minimum
supported version until June 2025.
Increase the minimum SciPy version to 1.11.1. (1.11.0 was yanked from PyPi because of licensing problems, so 1.11.1 is the oldest 1.11 release.) PiperOrigin-RevId: 713073731
2025-01-07 16:10:14 -08:00
* The minimum SciPy version is now 1.11. SciPy 1.11 will remain the minimum
supported version until June 2025.
jnp.einsum: default to optimize='auto'
2025-01-06 08:21:59 -08:00
* {func}`jax.numpy.einsum` now defaults to `optimize='auto'` rather than
`optimize='optimal'`. This avoids exponentially-scaling trace-time in
the case of many arguments ({jax-issue}`#25214`).
jnp.linalg.solve: finalize deprecation of batched 1D solves
2025-01-08 08:57:56 -08:00
* {func}`jax.numpy.linalg.solve` no longer supports batched 1D arguments
on the right hand side. To recover the previous behavior in these cases,
use `solve(a, b[..., None]).squeeze(-1)`.
Increase the minimum NumPy version to v1.25. Per SPEC 0, we drop NumPy v1.24 support on Dec 18, 2024.
2024-12-18 08:18:57 -05:00
Add support for N-D FFTs with D>3.
2024-12-19 12:02:36 +00:00
* New Features
* {func}`jax.numpy.fft.fftn`, {func}`jax.numpy.fft.rfftn`,
{func}`jax.numpy.fft.ifftn`, and {func}`jax.numpy.fft.irfftn` now support
transforms in more than 3 dimensions, which was previously the limit. See
{jax-issue}`#25606` for more details.
Add a register_custom_type_id function to the GPU plugins. This enables dynamic registration of custom FFI types on the appropriate platform via PJRT. PiperOrigin-RevId: 712904085
2025-01-07 07:28:21 -08:00
* Support added for user defined state in the FFI via the new
{func}`jax.ffi.register_ffi_type_id` function.
[aot] Add support for as_text(debug_info=True). This exposes an easier way to get StableHLO and HLO with more debugging information (source locations for StableHLO and metadata for HLO).
2025-01-09 12:03:05 +02:00
* The AOT lowering `.as_text()` method now supports the `debug_info` option
to include debugging information, e.g., source location, in the output.
Release JAX 0.5.0
2025-01-17 10:28:03 -05:00
changelog: link to api compatibility & python version docs
2024-12-23 09:39:45 -08:00
* Deprecations
* From {mod}`jax.interpreters.xla`, `abstractify` and `pytype_aval_mappings`
are now deprecated, having been replaced by symbols of the same name
in {mod}`jax.core`.
Deprecate scipy.special.lpmn & lpmn_values
2025-01-06 09:31:15 -08:00
* {func}`jax.scipy.special.lpmn` and {func}`jax.scipy.special.lpmn_values`
are deprecated, following their deprecation in SciPy v1.15.0. There are
no plans to replace these deprecated functions with new APIs.
Move jex.ffi to jax.ffi.
2024-12-20 11:26:04 +00:00
* The {mod}`jax.extend.ffi` submodule was moved to {mod}`jax.ffi`, and the
previous import path is deprecated.
changelog: link to api compatibility & python version docs
2024-12-23 09:39:45 -08:00
* Deletions
* `jax_enable_memories` flag has been deleted and the behavior of that flag
is on by default.
Finalize deprecation of some symbols from jax.lib.xla_client
2024-12-23 09:23:06 -08:00
* From `jax.lib.xla_client`, the previously-deprecated `Device` and
`XlaRuntimeError` symbols have been removed; instead use `jax.Device`
and `jax.errors.JaxRuntimeError` respectively.
Remove deprecated jax.experimental.array_api
2025-01-06 15:19:02 -08:00
* The `jax.experimental.array_api` module has been removed after being
deprecated in JAX v0.4.32. Since that release, {mod}`jax.numpy` supports
the array API directly.
changelog: link to api compatibility & python version docs
2024-12-23 09:39:45 -08:00
Update version numbers after release.
2024-12-17 18:03:28 -05:00
## jax 0.4.38 (Dec 17, 2024)
Merge branch 'release/0.4.37' into main.
2024-12-09 20:21:43 -05:00
Jax: Stop returning a list of cost-analyses. As it stands, there is only ever one element in this list (see b/384741132) and only the 0th element is ever used so we can simplify. This is a potentially breaking change for external users, but (as stated in the [documentation](https://jax.readthedocs.io/en/latest/aot.html#debug-information-and-analyses-when-available)) no guarantees are made on this type, which is intended for debugging purposes and not intended to be a reliable public API. PiperOrigin-RevId: 715837855
2025-01-15 09:53:29 -08:00
* Breaking Changes
* `XlaExecutable.cost_analysis` now returns a `dict[str, float]` (instead of a
single-element `list[dict[str, float]]`).
Add jax.tree shortcuts for .*_with_path calls, for convenience of users. PiperOrigin-RevId: 705645570
2024-12-12 15:13:01 -08:00
* Changes:
* `jax.tree.flatten_with_path` and `jax.tree.map_with_path` are added
as shortcuts of the corresponding `tree_util` functions.
jax.core: deprecate a number of APIs
2024-12-10 11:11:32 -08:00
* Deprecations
temporarily un-deprecate several jax.core APIs. These were causing excessive log-spam for some users; I'll work to migrate them to jax.extend before re-deprecating these.
2024-12-12 13:15:58 -08:00
* a number of APIs in the internal `jax.core` namespace have been deprecated.
Most were no-ops, were little-used, or can be replaced by APIs of the same
name in {mod}`jax.extend.core`; see the documentation for {mod}`jax.extend`
for information on the compatibility guarantees of these semi-public extensions.
Finalize some deprecations in jax.core, jax.lib.xla_bridge, and jax.lib.xla_client.
2024-12-11 09:50:33 -08:00
* Several previously-deprecated APIs have been removed, including:
* from {mod}`jax.core`: `check_eqn`, `check_type`, `check_valid_jaxtype`, and
`non_negative_dim`.
* from {mod}`jax.lib.xla_bridge`: `xla_client` and `default_backend`.
* from {mod}`jax.lib.xla_client`: `_xla` and `bfloat16`.
Finalize deprecation of jnp.round_ PiperOrigin-RevId: 705998500
2024-12-13 14:13:14 -08:00
* from {mod}`jax.numpy`: `round_`.
jax.core: deprecate a number of APIs
2024-12-10 11:11:32 -08:00
[export] Expand exporting to work with AbstractMesh. This is a follow up from #25640 that enabled lowering with AbstractMesh. This required adding `num_devices` to `lowering.compiler_args` because in presence of an AbstractMesh the device_assignment is not accurate.
2024-12-12 17:22:05 +02:00
* New Features
* {func}`jax.export.export` can be used for device-polymorphic export with
shardings constructed with {func}`jax.sharding.AbstractMesh`.
See the [jax.export documentation](https://jax.readthedocs.io/en/latest/export/export.html#device-polymorphic-export).
Reverts 525b646c0ebd5205f4fa0639c94adb2de47e1cf0 PiperOrigin-RevId: 707146329
2024-12-17 10:05:58 -08:00
* Added {func}`jax.lax.split`. This is a primitive version of
{func}`jax.numpy.split`, added because it yields a more compact
transpose during automatic differentiation.
[export] Expand exporting to work with AbstractMesh. This is a follow up from #25640 that enabled lowering with AbstractMesh. This required adding `num_devices` to `lowering.compiler_args` because in presence of an AbstractMesh the device_assignment is not accurate.
2024-12-12 17:22:05 +02:00
Merge branch 'release/0.4.37' into main.
2024-12-09 20:21:43 -05:00
## jax 0.4.37 (Dec 9, 2024)
Avoid index out of range error in carry structure check
2024-12-09 10:44:28 -08:00
Update versions for v0.4.37 release.
2024-12-09 15:38:37 -05:00
This is a patch release of jax 0.4.36. Only "jax" was released at this version.
Bump JAX version after release. PiperOrigin-RevId: 703472753
2024-12-06 05:57:30 -08:00
Avoid index out of range error in carry structure check
2024-12-09 10:44:28 -08:00
* Bug fixes
Update versions for v0.4.37 release.
2024-12-09 15:38:37 -05:00
* Fixed a bug where `jit` would error if an argument was named `f` (#25329).
Avoid index out of range error in carry structure check
2024-12-09 10:44:28 -08:00
* Fix a bug that will throw `index out of range` error in
{func}`jax.lax.while_loop` if the user register pytree node class with
different aux data for the flatten and flatten_with_path.
Update versions for v0.4.37 release.
2024-12-09 15:38:37 -05:00
* Pinned a new libtpu release (0.0.6) that fixes a compiler bug on TPU v6e.
Avoid index out of range error in carry structure check
2024-12-09 10:44:28 -08:00
Bump JAX version after release. PiperOrigin-RevId: 703472753
2024-12-06 05:57:30 -08:00
## jax 0.4.36 (Dec 5, 2024)
Merge branch 'release/0.4.35'
2024-10-22 17:02:22 -04:00
[jax2tf] Disable jax2tf with non-native serialization. jax2tf with native_serialization=False or with enable_xla=False have been deprecated since July 2024. This change turns an attempt to use `native_serialization=False` or `enable_xla=False` into an error. PiperOrigin-RevId: 689708392
2024-10-25 02:30:17 -07:00
* Breaking Changes
Mention stackless in the release notes.
2024-11-20 14:53:52 -05:00
* This release lands "stackless", an internal change to JAX's tracing
machinery. We made trace dispatch purely a function of context rather than a
function of both context and data. This let us delete a lot of machinery for
managing data-dependent tracing: levels, sublevels, `post_process_call`,
`new_base_main`, `custom_bind`, and so on. The change should only affect
users that use JAX internals.
If you do use JAX internals then you may need to
update your code (see
https://github.com/jax-ml/jax/commit/c36e1f7c1ad4782060cbc8e8c596d85dfb83986f
for clues about how to do this). There might also be version skew
issues with JAX libraries that do this. If you find this change breaks your
non-JAX-internals-using code then try the
`config.jax_data_dependent_tracing_fallback` flag as a workaround, and if
you need help updating your code then please file a bug.
[jax2tf] Disable jax2tf with non-native serialization. jax2tf with native_serialization=False or with enable_xla=False have been deprecated since July 2024. This change turns an attempt to use `native_serialization=False` or `enable_xla=False` into an error. PiperOrigin-RevId: 689708392
2024-10-25 02:30:17 -07:00
* {func}`jax.experimental.jax2tf.convert` with `native_serialization=False`
or with `enable_xla=False` have been deprecated since July 2024, with
JAX version 0.4.31. Now we removed support for these use cases. `jax2tf`
with native serialization will still be supported.
Finalize deprecation of xb, xc, & xe symbols in jax.interpreters.xla PiperOrigin-RevId: 689792265
2024-10-25 08:10:29 -07:00
* In `jax.interpreters.xla`, the `xb`, `xc`, and `xe` symbols have been removed
after being deprecated in JAX v0.4.31. Instead use `xb = jax.lib.xla_bridge`,
`xc = jax.lib.xla_client`, and `xe = jax.lib.xla_extension`.
Remove deprecated jax.experimental.export module. These tools are now available at jax.export.
2024-10-30 05:27:29 -07:00
* The deprecated module `jax.experimental.export` has been removed. It was replaced
by {mod}`jax.export` in JAX v0.4.30. See the [migration guide](https://jax.readthedocs.io/en/latest/export/export.html#migration-guide-from-jax-experimental-export)
for information on migrating to the new API.
Remove the `initial` argument to `jax.nn.softmax` and `jax.nn.log_softmax`. This argument was deprecated in JAX v0.4.27 and has no effect in JAX v0.4.27 and later. PiperOrigin-RevId: 693023366
2024-11-04 10:54:41 -08:00
* The `initial` argument to {func}`jax.nn.softmax` and {func}`jax.nn.log_softmax`
has been removed, after being deprecated in v0.4.27.
Error on numpy array conversion of PRNG key array
2024-11-05 09:26:42 -08:00
* Calling `np.asarray` on typed PRNG keys (i.e. keys produced by :func:`jax.random.key`)
now raises an error. Previously, this returned a scalar object array.
[export] Cleanup in the export module. With jax.experimental.export gone we can now do some cleanup in the export module. In particular we remove the `export.args_spec` API, and the `lowering_platforms` arg for `export.export`. These were deprecated in June 2024. PiperOrigin-RevId: 692398132
2024-11-01 22:55:56 -07:00
* The following deprecated methods and functions in {mod}`jax.export` have
been removed:
* `jax.export.DisabledSafetyCheck.shape_assertions`: it had no effect
already.
* `jax.export.Exported.lowering_platforms`: use `platforms`.
* `jax.export.Exported.mlir_module_serialization_version`:
use `calling_convention_version`.
* `jax.export.Exported.uses_shape_polymorphism`:
use `uses_global_constants`.
* the `lowering_platforms` kwarg for {func}`jax.export.export`: use
`platforms` instead.
[shape_poly] Remove some deprecated kwargs PiperOrigin-RevId: 703116755
2024-12-05 08:01:47 -08:00
* The kwargs `symbolic_scope` and `symbolic_constraints` from
{func}`jax.export.symbolic_args_specs` have been removed. They were
deprecated in June 2024. Use `scope` and `constraints` instead.
Make Tracers non-hashable
2024-11-05 09:08:33 -08:00
* Hashing of tracers, which has been deprecated since version 0.4.30, now
results in a `TypeError`.
Re-factor build CLI to a subcommand based approach This commit reworks the JAX build CLI to a subcommand based approach where CLI use cases are now defined as subcommands. Two subcommands are defined: build and requirements_update. "build" is to be used when wanting to build a JAX wheel package. "requirements_update" is to be used when wanting to update the requirements_lock.txt files. The new structure offers a clear and organized CLI that enables users to execute specific build tasks without having to navigate through a monolithic script. Each subcommand has specific arguments that apply to its respective build process. In addition, arguments are separated into groups to achieve a cleaner separation and improves the readability when the CLI subcommands are run with `--help`. It also makes it clear as to which parts of the build they affect. E.g: CUDA arguments only apply to CUDA builds, ROCM arguments only apply to ROCM builds, etc. This reduces the complexity and the potential for errors during the build process. Segregating functionalities into distinct subcommands also simplifies the code which should help with the maintenance and future extensions. There is also a transition from using `subprocess.check_output` to `asyncio.create_subprocess_shell` for executing the build commands which allows for streaming logs and helps in showing the build progress in real time. Usage: * Building `jaxlib`: ``` python build/build.py build --wheels=jaxlib --python_version=3.10 ``` * Building `jax-cuda-plugin`: ``` python build/build.py build --wheels=jax-cuda-plugin --cuda_version=12.3.2 --cudnn_version=9.1.1 --python_version=3.10 ``` * Building multiple packages: ``` python build/build.py build --wheels=jaxlib,jax-cuda-plugin,jax-cuda-pjrt --cuda_version=12.3.2 --cudnn_version=9.1.1 --python_version=3.10 ``` * Building `jax-rocm-pjrt`: ``` python build/build.py build --wheels=jax-rocm-pjrt --rocm_version=60 --rocm_path=/path/to/rocm ``` * Using a local XLA path: ``` python build/build.py build --wheels=jaxlib --local_xla_path=/path/to/xla ``` * Updating requirements_lock.txt files: ``` python build/build.py requirements_update --python_version=3.10 ``` For more details on each argument and to see available options, run: ``` python build/build.py build --help ``` or ``` python build/build.py requirements_update --help ``` PiperOrigin-RevId: 700075411
2024-11-25 13:02:22 -08:00
* Refactor: JAX build CLI (build/build.py) now uses a subcommand structure and
replaces previous build.py usage. Run `python build/build.py --help` for
more details. Brief overview of the new subcommand options:
* `build`: Builds JAX wheel packages. For e.g., `python build/build.py build --wheels=jaxlib,jax-cuda-pjrt`
* `requirements_update`: Updates requirements_lock.txt files.
jax.scipy.linalg.toeplitz: support implicit batching
2024-11-11 15:32:43 -08:00
* {func}`jax.scipy.linalg.toeplitz` now does implicit batching on multi-dimensional
inputs. To recover the previous behavior, you can call {func}`jax.numpy.ravel`
on the function inputs.
Update jax.scipy.special.gamma and gammasgn to return NaN for negative integer inputs. Change to match upstream scipy: https://github.com/scipy/scipy/pull/21827. Fixes #24875
2024-11-18 09:46:22 -05:00
* {func}`jax.scipy.special.gamma` and {func}`jax.scipy.special.gammasgn` now
return NaN for negative integer inputs, to match the behavior of SciPy from
https://github.com/scipy/scipy/pull/21827.
finalize deprecation of jax.clear_backends
2024-11-12 12:26:34 -08:00
* `jax.clear_backends` was removed after being deprecated in v0.4.26.
[export] Removed __gpu$xla.gpu.triton (Pallas GPU) from the list of custom calls with guaranteed compatibility. This is because the underlying Triton IR does not guarantee compatibility. PiperOrigin-RevId: 703127711
2024-12-05 08:39:48 -08:00
* We removed the custom call "__gpu$xla.gpu.triton" from the list of custom
call that we guarantee export stability. This is because this custom call
relies on Triton IR, which is not guaranteed to be stable. If you need
to export code that uses this custom call, you can use the `disabled_checks`
parameter. See more details in the [documentation](https://jax.readthedocs.io/en/latest/export/export.html#compatibility-guarantees-for-custom-calls).
[jax2tf] Disable jax2tf with non-native serialization. jax2tf with native_serialization=False or with enable_xla=False have been deprecated since July 2024. This change turns an attempt to use `native_serialization=False` or `enable_xla=False` into an error. PiperOrigin-RevId: 689708392
2024-10-25 02:30:17 -07:00
support exec_time_optimization_effort and memory_fitting_effort xla compilation options PiperOrigin-RevId: 692322944
2024-11-01 16:24:42 -07:00
* New Features
* {func}`jax.jit` got a new `compiler_options: dict[str, Any]` argument, for
passing compilation options to XLA. For the moment it's undocumented and
may be in flux.
register_dataclass: allow marking static fields via field(static=True)
2024-11-06 11:18:11 -08:00
* {func}`jax.tree_util.register_dataclass` now allows metadata fields to be
declared inline via {func}`dataclasses.field`. See the function documentation
for examples.
Add numpy.put_along_axis.
2024-11-14 15:23:26 -05:00
* Added {func}`jax.numpy.put_along_axis`.
Add a GPU implementation of `lax.linalg.eig`. This feature has been in the queue for a long time (see https://github.com/jax-ml/jax/issues/1259), and some folks have found that they can use `pure_callback` to call the CPU version as a workaround. It has recently come up that there can be issues when using `pure_callback` with JAX calls in the body (https://github.com/jax-ml/jax/issues/24255; this should be investigated separately). This change adds a native solution for computing `lax.linalg.eig` on GPU. By default, this is implemented by calling LAPACK on host directly because this has good performance for small to moderately sized problems (less than about 2048^2). For larger matrices, a GPU-backed implementation based on [MAGMA](https://icl.utk.edu/magma/) can have significantly better performance. (I should note that I haven't done a huge amount of benchmarking yet, but this was the breakeven point used by PyTorch, and I find roughly similar behavior so far.) We don't want to add MAGMA as a required dependency, but if a user has installed it, JAX can use it when the `jax_gpu_use_magma` configuration variable is set to `"on"`. By default, we try to dlopen `libmagma.so`, but the path to a non-standard installation location can be specified using the `JAX_GPU_MAGMA_PATH` environment variable. PiperOrigin-RevId: 697631402
2024-11-18 08:11:04 -08:00
* {func}`jax.lax.linalg.eig` and the related `jax.numpy` functions
({func}`jax.numpy.linalg.eig` and {func}`jax.numpy.linalg.eigvals`) are now
supported on GPU. See {jax-issue}`#24663` for more details.
Add exec_time_optimization_effort and memory_fitting_effort flags. These flags control the amount of effort the compiler spends minimizing execution time and memory usage, respectively. They can be set via the command line, e.g. . Valid values are between -1.0 and 1.0, default is 0.0.
2024-11-26 13:57:47 +00:00
* Added two new configuration flags, `jax_exec_time_optimization_effort` and `jax_memory_fitting_effort`, to control the amount of effort the compiler spends minimizing execution time and memory usage, respectively. Valid values are between -1.0 and 1.0, default is 0.0.
support exec_time_optimization_effort and memory_fitting_effort xla compilation options PiperOrigin-RevId: 692322944
2024-11-01 16:24:42 -07:00
Add GPU overflow bugfix (#24846) to changelog.
2024-11-12 08:57:52 -08:00
* Bug fixes
* Fixed a bug where the GPU implementations of LU and QR decomposition would
result in an indexing overflow for batch sizes close to int32 max. See
{jax-issue}`#24843` for more details.
Deprecate several private APIs in jax.lib
2024-11-20 08:48:26 -08:00
* Deprecations
* `jax.lib.xla_extension.ArrayImpl` and `jax.lib.xla_client.ArrayImpl` are deprecated;
use `jax.Array` instead.
* `jax.lib.xla_extension.XlaRuntimeError` is deprecated; use `jax.errors.JaxRuntimeError`
instead.
Merge branch 'release/0.4.35'
2024-10-22 17:02:22 -04:00
## jax 0.4.35 (Oct 22, 2024)
Merge branch 'release/0.4.34'
2024-10-04 10:53:28 -04:00
Simplify definition of jnp.isscalar The new semantics are to return True for any array-like object with zero dimensions. Previously we only returned True for zero-dimensional array-like objects with a weak type. This ends up being more confusing/suprising than it needs to be, and the weak type dependence is rarely useful in practice. PiperOrigin-RevId: 682656411
2024-10-05 07:11:44 -07:00
* Breaking Changes
* {func}`jax.numpy.isscalar` now returns True for any array-like object with
zero dimensions. Previously it only returned True for zero-dimensional
array-like objects with a weak dtype.
[host_callback] Remove most of the jax.experimental.host_callback module These APIs have been deprecated since March 2024 and they are subsumed by the new JAX external callbacks. See https://github.com/google/jax/issues/20385 for a discussion. PiperOrigin-RevId: 682830525
2024-10-06 01:09:56 -07:00
* `jax.experimental.host_callback` has been deprecated since March 2024, with
JAX version 0.4.26. Now we removed it.
See {jax-issue}`#20385` for a discussion of alternatives.
Simplify definition of jnp.isscalar The new semantics are to return True for any array-like object with zero dimensions. Previously we only returned True for zero-dimensional array-like objects with a weak type. This ends up being more confusing/suprising than it needs to be, and the weak type dependence is rarely useful in practice. PiperOrigin-RevId: 682656411
2024-10-05 07:11:44 -07:00
Add lax.FftType. We had never provided a public name for the enum of FFT types; instead it was only known by a semi-private name (jax.lib.xla_client.FftType). Add a public name (jax.lax.FftType) and deprecate the private one. We define a new FftType IntEnum rather than trying to expose the one in xla_client. The xla_client definition was useful when building classic HLO, but we no longer do that so there's no reason we need to couple our type to XLA's type. PiperOrigin-RevId: 684447186
2024-10-10 08:06:50 -07:00
* Changes:
* `jax.lax.FftType` was introduced as a public name for the enum of FFT
operations. The semi-public API `jax.lib.xla_client.FftType` has been
deprecated.
Use libtpu releases rather than libtpu-nightly for jax[tpu]. PiperOrigin-RevId: 688632409
2024-10-22 11:45:58 -07:00
* TPU: JAX now installs TPU support from the `libtpu` package rather than
`libtpu-nightly`. For the next few releases JAX will pin an empty version of
`libtpu-nightly` as well as `libtpu` to ease the transition; that dependency
will be removed in Q1 2025.
Add lax.FftType. We had never provided a public name for the enum of FFT types; instead it was only known by a semi-private name (jax.lib.xla_client.FftType). Add a public name (jax.lax.FftType) and deprecate the private one. We define a new FftType IntEnum rather than trying to expose the one in xla_client. The xla_client definition was useful when building classic HLO, but we no longer do that so there's no reason we need to couple our type to XLA's type. PiperOrigin-RevId: 684447186
2024-10-10 08:06:50 -07:00
Deprecate jax.lib.xla_client.PaddingType. This type is unused by JAX, so there is no replacement. (JAX does have an internal PaddingType enum in lax, but it is not present in any APIs, as best I can tell.) PiperOrigin-RevId: 684451556
2024-10-10 08:21:28 -07:00
* Deprecations:
* The semi-public API `jax.lib.xla_client.PaddingType` has been deprecated.
No JAX APIs consume this type, so there is no replacement.
Fix listing of `vectorized` deprecation in changelog. As noted in https://github.com/jax-ml/jax/pull/23881, that change didn't actually make it in in time for the v0.4.34 release so I've moved it to the v0.4.35 section.
2024-10-08 12:55:47 -04:00
* 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` for more details.
Deprecate a number of APIs in jax.lib.xla_client. (Technically these aren't public, so they don't need a deprecation period, but this is the polite thing to do.) PiperOrigin-RevId: 684906277
2024-10-11 11:42:08 -07:00
* The semi-public API `jax.lib.xla_client.register_custom_call_target` has
been deprecated. Use the JAX FFI instead.
* The semi-public APIs `jax.lib.xla_client.dtype_to_etype`,
Add GPU overflow bugfix (#24846) to changelog.
2024-11-12 08:57:52 -08:00
`jax.lib.xla_client.ops`,
Deprecate a number of APIs in jax.lib.xla_client. (Technically these aren't public, so they don't need a deprecation period, but this is the polite thing to do.) PiperOrigin-RevId: 684906277
2024-10-11 11:42:08 -07:00
`jax.lib.xla_client.shape_from_pyval`, `jax.lib.xla_client.PrimitiveType`,
`jax.lib.xla_client.Shape`, `jax.lib.xla_client.XlaBuilder`, and
`jax.lib.xla_client.XlaComputation` have been deprecated. Use StableHLO
instead.
Deprecate jax.lib.xla_client.PaddingType. This type is unused by JAX, so there is no replacement. (JAX does have an internal PaddingType enum in lax, but it is not present in any APIs, as best I can tell.) PiperOrigin-RevId: 684451556
2024-10-10 08:21:28 -07:00
Fix wrong date in changelog
2024-10-06 23:16:30 +08:00
## jax 0.4.34 (October 4, 2024)
Update version numbers after v0.4.32 release.
2024-09-11 16:08:47 -04:00
Add Python 3.13 wheels to changelog.
2024-09-18 18:57:03 +00:00
* New Functionality
* This release includes wheels for Python 3.13. Free-threading mode is not yet
supported.
Add jax.errors.JaxRuntimeError as a public alias for the XlaRuntimeError class. Deprecate jax.lib.xla_client.XlaRuntimeError, which is not a public API. PiperOrigin-RevId: 679163106
2024-09-26 08:38:46 -07:00
* `jax.errors.JaxRuntimeError` has been added as a public alias for the
formerly private `XlaRuntimeError` type.
Add Python 3.13 wheels to changelog.
2024-09-18 18:57:03 +00:00
Reverts dffac29e63de6a51047fe77cf9d553ab762ef19b PiperOrigin-RevId: 678748794
2024-09-25 10:13:53 -07:00
* 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.
[host_callback] Flip the JAX_HOST_CALLBACK_LEGACY flag to False `jax.experimental.host_callback` has been deprecated since March 2024 (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 https://github.com/google/jax/issues/20385 for a discussion. PiperOrigin-RevId: 681004255
2024-10-01 07:06:48 -07:00
* `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` for a discussion.
Reverts dffac29e63de6a51047fe77cf9d553ab762ef19b PiperOrigin-RevId: 678748794
2024-09-25 10:13:53 -07:00
Deprecate passing NdArrays with ndim != 1 and non-arraylike inputs to jnp.trim_zeros
2024-09-18 17:06:28 +05:30
* Deprecations
Add Python 3.13 wheels to changelog.
2024-09-18 18:57:03 +00:00
* 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.
Finalize deprecation of pretty-printing utils in `jax.core.pp_*` PiperOrigin-RevId: 678775782
2024-09-25 11:19:20 -07:00
* Internal pretty-printing tools `jax.core.pp_*` have been removed, after
being deprecated in JAX v0.4.30.
Deprecate jax.lib.xla_client.Device. jax.Device is a longstanding public name for this class. PiperOrigin-RevId: 679197718
2024-09-26 10:16:20 -07:00
* `jax.lib.xla_client.Device` is deprecated; use `jax.Device` instead.
Add jax.errors.JaxRuntimeError as a public alias for the XlaRuntimeError class. Deprecate jax.lib.xla_client.XlaRuntimeError, which is not a public API. PiperOrigin-RevId: 679163106
2024-09-26 08:38:46 -07:00
* `jax.lib.xla_client.XlaRuntimeError` has been deprecated. Use
`jax.errors.JaxRuntimeError` instead.
Deprecate passing NdArrays with ndim != 1 and non-arraylike inputs to jnp.trim_zeros
2024-09-18 17:06:28 +05:30
Reverts 255c30303d32e7473262b2e35348175c87e4348f PiperOrigin-RevId: 674083626
2024-09-12 18:13:51 -07:00
* Deletion:
Delete jax.xla_computation since it's been 3 months since it was deprecated. PiperOrigin-RevId: 673938336
2024-09-12 11:47:03 -07:00
* `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).
* For cross-backend lowering, you can replace
`jax.xla_computation(fn, backend='tpu')(*args, **kwargs)` with
`jax.jit(fn).trace(*args, **kwargs).lower(lowering_platforms=('tpu',)).compiler_ir('hlo')`.
Removed the `named_shape` argument from `jex.core.ShapedArray` and `jax.ShapeDtypeStruct` It is unused and was only kept around to avoid breaking internal users. PiperOrigin-RevId: 674310795
2024-09-13 08:37:32 -07:00
* {class}`jax.ShapeDtypeStruct` no longer accepts the `named_shape` argument.
The argument was only used by `xmap` which was removed in 0.4.31.
Reverts dffac29e63de6a51047fe77cf9d553ab762ef19b PiperOrigin-RevId: 678748794
2024-09-25 10:13:53 -07:00
* `jax.tree.map(f, None, non-None)`, which previously emitted a
`DeprecationWarning`, now raises an error in a future version of jax. `None`
is only a tree-prefix of itself. To preserve the current behavior, you can
ask `jax.tree.map` to treat `None` as a leaf value by writing:
`jax.tree.map(lambda x, y: None if x is None else f(x, y), a, b, is_leaf=lambda x: x is None)`.
Finalize deprecation of XLACompatibleSharding PiperOrigin-RevId: 681156145
2024-10-01 14:01:48 -07:00
* `jax.sharding.XLACompatibleSharding` has been removed. Please use
`jax.sharding.Sharding`.
Delete jax.xla_computation since it's been 3 months since it was deprecated. PiperOrigin-RevId: 673938336
2024-09-12 11:47:03 -07:00
Fix an incorrect output for jnp.cumsum. If dtype=bool but a non-bool input is passed, we should test for non-equality with zero rather than performing a cast to integer.
2024-09-24 13:13:29 +00:00
* Bug fixes
* Fixed a bug where {func}`jax.numpy.cumsum` would produce incorrect outputs
if a non-boolean input was provided and `dtype=bool` was specified.
Edit implementation of jax.numpy.ldexp to get correct gradient.
2024-09-30 18:27:39 -04:00
* Edit implementation of {func}`jax.numpy.ldexp` to get correct gradient.
Update version numbers after v0.4.32 release.
2024-09-11 16:08:47 -04:00
Merge release/0.4.33 into main and update version numbers.
2024-09-16 18:46:24 +00:00
## jax 0.4.33 (September 16, 2024)
Prepare for v0.4.33 release. This release is branched off the v0.4.32 release, with two changes: a) a fixed libtpu pin, and b) a patch to revert an F64 tanh issue on CPU.
2024-09-16 13:20:56 +00:00
This is a patch release on top of jax 0.4.32, that fixes two bugs found in that
release.
A TPU-only data corruption bug was found in the version of libtpu pinned by
JAX 0.4.32, which manifested only if multiple TPU slices were present in the
same job, for example, if training on multiple v5e slices.
This release fixes that issue by pinning a fixed version of `libtpu`.
This release fixes an inaccurate result for F64 tanh on CPU (#23590).
Update version numbers after v0.4.32 release.
2024-09-11 16:08:47 -04:00
## jax 0.4.32 (September 11, 2024)
Finish jax and jaxlib v0.4.31 release PiperOrigin-RevId: 657388782
2024-07-29 17:56:53 -07:00
Prepare for v0.4.33 release. This release is branched off the v0.4.32 release, with two changes: a) a fixed libtpu pin, and b) a patch to revert an F64 tanh issue on CPU.
2024-09-16 13:20:56 +00:00
Note: This release was yanked from PyPi because of a data corruption bug on TPU.
See the 0.4.33 release notes for more details.
Finish jax and jaxlib v0.4.31 release PiperOrigin-RevId: 657388782
2024-07-29 17:56:53 -07:00
Add the FFI functions and tutorial to the changelog. Although we soft launched the FFI with v0.4.31, it would be nice to include an update in the changelog to help with visibility.
2024-09-06 12:30:28 -04:00
* New Functionality
* Added {func}`jax.extend.ffi.ffi_call` and {func}`jax.extend.ffi.ffi_lowering`
to support the use of the new {ref}`ffi-tutorial` to interface with custom
C++ and CUDA code from JAX.
[array api] Finalize array API in jax.numpy & deprecate jax.experimental.array_api
2024-08-01 11:19:17 -07:00
* Changes
Set `jax_enable_memories` flag to `True` by default PiperOrigin-RevId: 660579462
2024-08-07 16:24:42 -07:00
* `jax_enable_memories` flag is set to `True` by default.
[array api] Finalize array API in jax.numpy & deprecate jax.experimental.array_api
2024-08-01 11:19:17 -07:00
* {mod}`jax.numpy` now supports v2023.12 of the Python Array API Standard.
See {ref}`python-array-api` for more information.
Async dispatch expensive computations on the JAX CPU backend. By setting `jax.config.update('jax_cpu_enable_async_dispatch', False)`, one could opt out of the change and recover the old behavior. PiperOrigin-RevId: 659741822
2024-08-05 17:47:34 -07:00
* Computations on the CPU backend may now be dispatched asynchronously in
more cases. Previously non-parallel computations were always dispatched
synchronously. You can recover the old behavior by setting
`jax.config.update('jax_cpu_enable_async_dispatch', False)`.
Add a `jax.process_indices` function. The `jax.host_ids` function has be long deprecated, but the suggested alternative of `list(range(jax.process_count()))` relies on the current behavior that the list of process indices is always dense. In the future we may want to allow dynamic addition and removal of processes in which case `jax.process_count` and `jax.process_indices` would need to be updated, and it is useful for users to be able to use this forward-compatible interface. PiperOrigin-RevId: 662142636
2024-08-12 10:29:15 -07:00
* Added new {func}`jax.process_indices` function to replace the
`jax.host_ids()` function that was deprecated in JAX v0.2.13.
Update jnp.fabs to emulate the behavior of np.fabs for complex inputs
2024-08-28 20:16:09 +05:30
* To align with the behavior of `numpy.fabs`, `jax.numpy.fabs` has been
modified to no longer support `complex dtypes`.
``jax.tree_util.register_dataclass`` now validates ``data_fields`` and ``meta_fields`` A well-behaved registration call must list all ``init=True`` fields in either ``data_fields`` or ``meta_fields``. Otherwise, ``flatten . unflatten`` could potentially *not* be an identity PiperOrigin-RevId: 669244669
2024-08-30 02:01:02 -07:00
* ``jax.tree_util.register_dataclass`` now checks that ``data_fields``
and ``meta_fields`` includes all dataclass fields with ``init=True``
and only them, if ``nodetype`` is a dataclass.
First pass at ufunc interfaces for several jax.numpy functions
2024-08-30 11:53:02 -07:00
* Several {mod}`jax.numpy` functions now have full {class}`~jax.numpy.ufunc`
interfaces, including {obj}`~jax.numpy.add`, {obj}`~jax.numpy.multiply`,
{obj}`~jax.numpy.bitwise_and`, {obj}`~jax.numpy.bitwise_or`,
{obj}`~jax.numpy.bitwise_xor`, {obj}`~jax.numpy.logical_and`,
{obj}`~jax.numpy.logical_and`, and {obj}`~jax.numpy.logical_and`.
In future releases we plan to expand these to other ufuncs.
Make optimization_barrier a public lax API.
2024-09-05 19:49:12 +00:00
* Added {func}`jax.lax.optimization_barrier`, which allows users to prevent
compiler optimizations such as common-subexpression elimination and to
control scheduling.
[array api] Finalize array API in jax.numpy & deprecate jax.experimental.array_api
2024-08-01 11:19:17 -07:00
Drop support for mhlo in JAX's public API. PiperOrigin-RevId: 657551590
2024-07-30 05:27:23 -07:00
* Breaking changes
* The MHLO MLIR dialect (`jax.extend.mlir.mhlo`) has been removed. Use the
`stablehlo` dialect instead.
Raise ValueError for complex inputs to jnp.clip and jnp.hypot. Such inputs were deprecated in JAX v0.4.27, and have been raising a DeprecationWarning for the last several releases. PiperOrigin-RevId: 657717875
2024-07-30 13:48:51 -07:00
* Deprecations
* Complex inputs to {func}`jax.numpy.clip` and {func}`jax.numpy.hypot` are
no longer allowed, after being deprecated since JAX v0.4.27.
Deprecate several APIs in jax.lib.xla_bridge PiperOrigin-RevId: 658274719
2024-07-31 22:59:49 -07:00
* Deprecated the following APIs:
* `jax.lib.xla_bridge.xla_client`: use {mod}`jax.lib.xla_client` directly.
* `jax.lib.xla_bridge.get_backend`: use {func}`jax.extend.backend.get_backend`.
* `jax.lib.xla_bridge.default_backend`: use {func}`jax.extend.backend.default_backend`.
[array api] Finalize array API in jax.numpy & deprecate jax.experimental.array_api
2024-08-01 11:19:17 -07:00
* The `jax.experimental.array_api` module is deprecated, and importing it is no
longer required to use the Array API. `jax.numpy` supports the array API
directly; see {ref}`python-array-api` for more information.
Deprecate several internal utilities in jax.core
2024-08-14 10:06:13 -07:00
* The internal utilities `jax.core.check_eqn`, `jax.core.check_type`, and
`jax.core.check_valid_jaxtype` are now deprecated, and will be removed in
the future.
Deprecate jax.numpy.round_ NumPy removed np.round in version 2.0; jax.numpy.round is drop-in replacement.
2024-09-03 06:52:07 -07:00
* `jax.numpy.round_` has been deprecated, following removal of the corresponding
API in NumPy 2.0. Use {func}`jax.numpy.round` instead.
Deprecated calling ``jax.dlpack.from_dlpack`` with a DLPack tensor PiperOrigin-RevId: 670723176
2024-09-03 15:15:19 -07:00
* Passing a DLPack capsule to {func}`jax.dlpack.from_dlpack` is deprecated.
The argument to {func}`jax.dlpack.from_dlpack` should be an array from
another framework that implements the ``__dlpack__`` protocol.
Raise ValueError for complex inputs to jnp.clip and jnp.hypot. Such inputs were deprecated in JAX v0.4.27, and have been raising a DeprecationWarning for the last several releases. PiperOrigin-RevId: 657717875
2024-07-30 13:48:51 -07:00
Update version numbers after v0.4.32 release.
2024-09-11 16:08:47 -04:00
## jaxlib 0.4.32 (September 11, 2024)
Finish jax and jaxlib v0.4.31 release PiperOrigin-RevId: 657388782
2024-07-29 17:56:53 -07:00
Prepare for v0.4.33 release. This release is branched off the v0.4.32 release, with two changes: a) a fixed libtpu pin, and b) a patch to revert an F64 tanh issue on CPU.
2024-09-16 13:20:56 +00:00
Note: This release was yanked from PyPi because of a data corruption bug on TPU.
See the 0.4.33 release notes for more details.
Finish jax and jaxlib v0.4.31 release PiperOrigin-RevId: 657388782
2024-07-29 17:56:53 -07:00
Introduce hermetic CUDA in Google ML projects. 1) Hermetic CUDA rules allow building wheels with GPU support on a machine without GPUs, as well as running Bazel GPU tests on a machine with only GPUs and NVIDIA driver installed. When `--config=cuda` is provided in Bazel options, Bazel will download CUDA, CUDNN and NCCL redistributions in the cache, and use them during build and test phases. [Default location of CUNN redistributions](https://developer.download.nvidia.com/compute/cudnn/redist/) [Default location of CUDA redistributions](https://developer.download.nvidia.com/compute/cuda/redist/) [Default location of NCCL redistributions](https://pypi.org/project/nvidia-nccl-cu12/#history) 2) To include hermetic CUDA rules in your project, add the following in the WORKSPACE of the downstream project dependent on XLA. Note: use `@local_tsl` instead of `@tsl` in Tensorflow project. ``` load( "@tsl//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl", "cuda_json_init_repository", ) cuda_json_init_repository() load( "@cuda_redist_json//:distributions.bzl", "CUDA_REDISTRIBUTIONS", "CUDNN_REDISTRIBUTIONS", ) load( "@tsl//third_party/gpus/cuda/hermetic:cuda_redist_init_repositories.bzl", "cuda_redist_init_repositories", "cudnn_redist_init_repository", ) cuda_redist_init_repositories( cuda_redistributions = CUDA_REDISTRIBUTIONS, ) cudnn_redist_init_repository( cudnn_redistributions = CUDNN_REDISTRIBUTIONS, ) load( "@tsl//third_party/gpus/cuda/hermetic:cuda_configure.bzl", "cuda_configure", ) cuda_configure(name = "local_config_cuda") load( "@tsl//third_party/nccl/hermetic:nccl_redist_init_repository.bzl", "nccl_redist_init_repository", ) nccl_redist_init_repository() load( "@tsl//third_party/nccl/hermetic:nccl_configure.bzl", "nccl_configure", ) nccl_configure(name = "local_config_nccl") ``` PiperOrigin-RevId: 662981325
2024-08-14 10:57:53 -07:00
* Breaking changes
Add a note to the changelog about the new CPU thunks backend, enabled in 0.4.32.
2024-09-19 15:37:52 -04:00
* This release of jaxlib switched to a new version of the CPU backend, which
should compile faster and leverage parallelism better. If you experience
any problems due to this change, you can temporarily enable the old CPU
backend by setting the environment variable
`XLA_FLAGS=--xla_cpu_use_thunk_runtime=false`. If you need to do this,
please file a JAX bug with instructions to reproduce.
Introduce hermetic CUDA in Google ML projects. 1) Hermetic CUDA rules allow building wheels with GPU support on a machine without GPUs, as well as running Bazel GPU tests on a machine with only GPUs and NVIDIA driver installed. When `--config=cuda` is provided in Bazel options, Bazel will download CUDA, CUDNN and NCCL redistributions in the cache, and use them during build and test phases. [Default location of CUNN redistributions](https://developer.download.nvidia.com/compute/cudnn/redist/) [Default location of CUDA redistributions](https://developer.download.nvidia.com/compute/cuda/redist/) [Default location of NCCL redistributions](https://pypi.org/project/nvidia-nccl-cu12/#history) 2) To include hermetic CUDA rules in your project, add the following in the WORKSPACE of the downstream project dependent on XLA. Note: use `@local_tsl` instead of `@tsl` in Tensorflow project. ``` load( "@tsl//third_party/gpus/cuda/hermetic:cuda_json_init_repository.bzl", "cuda_json_init_repository", ) cuda_json_init_repository() load( "@cuda_redist_json//:distributions.bzl", "CUDA_REDISTRIBUTIONS", "CUDNN_REDISTRIBUTIONS", ) load( "@tsl//third_party/gpus/cuda/hermetic:cuda_redist_init_repositories.bzl", "cuda_redist_init_repositories", "cudnn_redist_init_repository", ) cuda_redist_init_repositories( cuda_redistributions = CUDA_REDISTRIBUTIONS, ) cudnn_redist_init_repository( cudnn_redistributions = CUDNN_REDISTRIBUTIONS, ) load( "@tsl//third_party/gpus/cuda/hermetic:cuda_configure.bzl", "cuda_configure", ) cuda_configure(name = "local_config_cuda") load( "@tsl//third_party/nccl/hermetic:nccl_redist_init_repository.bzl", "nccl_redist_init_repository", ) nccl_redist_init_repository() load( "@tsl//third_party/nccl/hermetic:nccl_configure.bzl", "nccl_configure", ) nccl_configure(name = "local_config_nccl") ``` PiperOrigin-RevId: 662981325
2024-08-14 10:57:53 -07:00
* Hermetic CUDA support is added.
Hermetic CUDA uses a specific downloadable version of CUDA instead of the
user’s locally installed CUDA. Bazel will download CUDA, CUDNN and NCCL
distributions, and then use CUDA libraries and tools as dependencies in
various Bazel targets. This enables more reproducible builds for JAX and its
supported CUDA versions.
Add a callout that LibTPU now supports profiling of SparseCore for TPUv5p chips which will be viewable in Tensorboard Profiler's TraceViewer tool. PiperOrigin-RevId: 667708094
2024-08-26 14:04:07 -07:00
* Changes
* SparseCore profiling is added.
* JAX now supports profiling [SparseCore](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#sparsecore) on TPUv5p chips. These traces will be viewable in Tensorboard Profiler's [TraceViewer](https://www.tensorflow.org/guide/profiler#trace_viewer).
Finish jax and jaxlib v0.4.31 release PiperOrigin-RevId: 657388782
2024-07-29 17:56:53 -07:00
## jax 0.4.31 (July 29, 2024)
Delete `xmap` and the `jax.experimental.maps` module. It's been 5 months since its deprecation (more than the standard 3 months deprecation period). PiperOrigin-RevId: 655614395
2024-07-24 10:23:29 -07:00
* Deletion
* xmap has been deleted. Please use {func}`shard_map` as the replacement.
Finish jax and jaxlib 0.4.30 release PiperOrigin-RevId: 644402635
2024-06-18 08:55:26 -07:00
Update minimum NumPy version to v1.24.
2024-06-21 14:57:35 -07:00
* Changes
Bump the minimum CUDNN version to v9.1. This actually was already the minimum version since we build with that version, but we needed to tighten the constraints. Also in passing, drop mentions of CUDA builds from the Windows build instructions. jaxlib hasn't built with CUDA enabled on Windows for a very long time, so it's probably best we just don't mention it. PiperOrigin-RevId: 657225917
2024-07-29 09:27:54 -07:00
* The minimum CuDNN version is v9.1. This was true in previous releases also,
but we now declare this version constraint formally.
Update minimum Python version to 3.10.
2024-06-26 13:43:15 -04:00
* The minimum Python version is now 3.10. 3.10 will remain the minimum
supported version until July 2025.
* The minimum NumPy version is now 1.24. NumPy 1.24 will remain the minimum
supported version until December 2024.
Bump minimum SciPy version to 1.10. SciPy 1.9.0 was released July 29, 2022, which is 24 months ago PiperOrigin-RevId: 657215038
2024-07-29 08:49:41 -07:00
* The minimum SciPy version is now 1.10. SciPy 1.10 will remain the minimum
supported version until January 2025.
Updated jnp.ceil/floor/trunc to preserve int dtypes Description: - Updated jnp.ceil/floor/trunc to preserve int dtypes - Updated tests - For integral dtypes but we can't yet today compare types vs numpy as numpy 2.0.0rc2 is not yet array api compliant in this case
2024-05-27 10:51:13 +00:00
* {func}`jax.numpy.ceil`, {func}`jax.numpy.floor` and {func}`jax.numpy.trunc` now return the output
of the same dtype as the input, i.e. no longer upcast integer or boolean inputs to floating point.
`libdevice.10.bc` is removed from JAX wheels bundle. The recommended source of JAX wheels is `pip`, and NVIDIA dependencies are installed automatically when JAX is installed via `pip install`. `libdevice` gets installed from `nvidia-cuda-nvcc-cu12` package. PiperOrigin-RevId: 647328834
2024-06-27 08:35:09 -07:00
* `libdevice.10.bc` is no longer bundled with CUDA wheels. It must be
installed either as a part of local CUDA installation, or via NVIDIA's CUDA
pip wheels.
Changed ``pl.BlockSpec`` to accept ``block_shape`` before ``index_map`` So, instead of pl.BlockSpec(lambda i, j: ..., (42, 24)) ``pl.BlockSpec`` now expects pl.BlockSpec((42, 24), lambda i, j: ...) I will update Pallas tests in a follow up. PiperOrigin-RevId: 648486321
2024-07-01 14:25:25 -07:00
* {class}`jax.experimental.pallas.BlockSpec` now expects `block_shape` to
be passed *before* `index_map`. The old argument order is deprecated and
will be removed in a future release.
Updated the repr of GPU devices to be more consistent with TPUs/CPUs. For example, `cuda(id=0)` will now be `CudaDevice(id=0)` PiperOrigin-RevId: 651393690
2024-07-11 06:53:32 -07:00
* Updated the repr of gpu devices to be more consistent
with TPUs/CPUs. For example, `cuda(id=0)` will now be `CudaDevice(id=0)`.
[array API] add device property & to_device method
2024-07-23 09:48:51 -07:00
* Added the `device` property and `to_device` method to {class}`jax.Array`, as
part of JAX's [Array API](https://data-apis.org/array-api) support.
Finalize deprecation of several previously-deprecated jax.core functions: - `jax.core.canonicalize_shape` - `jax.core.dimension_as_value` - `jax.core.definitely_equal` - `jax.core.symbolic_equal_dim` These have been raising deprecation warnings since JAX v0.4.24, released Feb 6 2024. PiperOrigin-RevId: 647671122
2024-06-28 07:27:23 -07:00
* Deprecations
* Removed a number of previously-deprecated internal APIs related to
polymorphic shapes. From {mod}`jax.core`: removed `canonicalize_shape`,
`dimension_as_value`, `definitely_equal`, and `symbolic_equal_dim`.
Deprecate support for custom lowering rules that return tuple-wrapped ir.Values. https://github.com/google/jax/pull/22211 forbade custom lowering rules from returning singleton tuples of ir.Value, but this appears to break downstream users, notably Transformer Engine. Instead, allow lowering rules to return singleton tuples and unwrap them if needed, but warn if this behavior is seen. PiperOrigin-RevId: 650345051
2024-07-08 12:54:04 -07:00
* HLO lowering rules should no longer wrap singleton ir.Values in tuples.
Instead, return singleton ir.Values unwrapped. Support for wrapped values
will be removed in a future version of JAX.
[jax2tf] Deprecate jax2tf with native_serialization=False or enable_xla=False. Also disable many of the non-native-serialization jax2tf tests. In particular, I am disabling the thousands of primitives tests in graph serialization mode. I kept jax2tf_test running in both native and graph serialization mode. PiperOrigin-RevId: 652749891
2024-07-16 02:04:59 -07:00
* {func}`jax.experimental.jax2tf.convert` with `native_serialization=False`
or `enable_xla=False` is now deprecated and this support will be removed in
a future version.
Native serialization has been the default since JAX 0.4.16 (September 2023).
Finalize deprecation of jax.random.shuffle This has been raising a DeprecationWarning for longer than anyone can remember. PiperOrigin-RevId: 656765001
2024-07-27 11:21:02 -07:00
* The previously-deprecated function `jax.random.shuffle` has been removed;
instead use `jax.random.permutation` with `independent=True`.
Update minimum NumPy version to v1.24.
2024-06-21 14:57:35 -07:00
Finish jax and jaxlib v0.4.31 release PiperOrigin-RevId: 657388782
2024-07-29 17:56:53 -07:00
## jaxlib 0.4.31 (July 29, 2024)
Finish jax and jaxlib 0.4.30 release PiperOrigin-RevId: 644402635
2024-06-18 08:55:26 -07:00
[JAX] Teach jit fast path how to handle negative static_argnums correctly. PiperOrigin-RevId: 645172085
2024-06-20 15:17:37 -07:00
* Bug fixes
* Fixed a bug that meant that negative static_argnums to a jit were mishandled
by the jit dispatch fast path.
Don't mask out zero elements on the diagonal of the matrix when inverting triangular matrices. The intent of this code seems to have been to mask out zeros that were part of padding on the diagonal. However, this isn't correct: if there is a zero on the diagonal, we very much want to get an inf or nan! We also appear to now pad with the identity matrix. Fixes https://github.com/google/jax/issues/3589 Fixes https://github.com/google/jax/issues/15429 PiperOrigin-RevId: 653562611
2024-07-18 04:08:55 -07:00
* Fixed a bug that meant triangular solves of batches of singular matrices
produce nonsensical finite values, instead of inf or nan (#3589, #15429).
[JAX] Teach jit fast path how to handle negative static_argnums correctly. PiperOrigin-RevId: 645172085
2024-06-20 15:17:37 -07:00
Finish jax and jaxlib 0.4.30 release PiperOrigin-RevId: 644402635
2024-06-18 08:55:26 -07:00
## jax 0.4.30 (June 18, 2024)
Update version numbers after v0.4.29 release.
2024-06-10 14:33:13 -04:00
[Rollback] Bumped the minimum ml_dtypes version to 0.4.0 Reverts e86c436e7f8e4e0546eff8bc2d3756a7c49dc83b PiperOrigin-RevId: 642741832
2024-06-12 14:39:11 -07:00
* Changes
* JAX supports ml_dtypes >= 0.2. In 0.4.29 release, the ml_dtypes version was
bumped to 0.4.0 but this has been rolled back in this release to give users
of both TensorFlow and JAX more time to migrate to a newer TensorFlow
release.
Add registration handler for TPU v5e in mesh_utils. PiperOrigin-RevId: 643092629
2024-06-13 12:51:43 -07:00
* `jax.experimental.mesh_utils` can now create an efficient mesh for TPU v5e.
Update JAX dependencies, extras, and documentation for plugins. * Make jaxlib a direct dependency of jax. * Remove mentions of monolithic CUDA installations from the JAX documentation. * Drop the cuda12_pip extra and the cudnn version specific extras. * Add a with_cuda extra to the jax-cuda12-plugin package, use it in jax's setup.py. This allows us to specify cuda extras in one place. * Make a few small doc improvements.
2024-06-13 10:50:24 -04:00
* jax now depends on jaxlib directly. This change was enabled by the CUDA
plugin switch: there are no longer multiple jaxlib variants. You can install
a CPU-only jax with `pip install jax`, no extras required.
Re-land https://github.com/google/jax/pull/21847 Reverts 0bcc81ceb33e3065110e3dd56ca215dbb62f0a7b PiperOrigin-RevId: 643202512
2024-06-13 19:53:10 -07:00
* Added an API for exporting and serializing JAX functions. This used
to exist in `jax.experimental.export` (which is being deprecated),
and will now live in `jax.export`.
See the [documentation](https://jax.readthedocs.io/en/latest/export/index.html).
Add registration handler for TPU v5e in mesh_utils. PiperOrigin-RevId: 643092629
2024-06-13 12:51:43 -07:00
Deprecate internal pretty-printing APIs, jax.core.pp_*
2024-06-13 09:31:43 -07:00
* Deprecations
* Internal pretty-printing tools `jax.core.pp_*` are deprecated, and will be removed
in a future release.
Deprecate hashing of tracers
2024-06-13 13:14:27 -07:00
* Hashing of tracers is deprecated, and will lead to a `TypeError` in a future JAX
release. This previously was the case, but there was an inadvertent regression in
the last several JAX releases.
Re-land https://github.com/google/jax/pull/21847 Reverts 0bcc81ceb33e3065110e3dd56ca215dbb62f0a7b PiperOrigin-RevId: 643202512
2024-06-13 19:53:10 -07:00
* `jax.experimental.export` is deprecated. Use {mod}`jax.export` instead.
See the [migration guide](https://jax.readthedocs.io/en/latest/export/export.html#migration-guide-from-jax-experimental-export).
Deprecate passing of arrays in place of dtypes.
2024-06-13 10:56:12 -07:00
* Passing an array in place of a dtype is now deprecated in most cases; e.g. for arrays
`x` and `y`, `x.astype(y)` will raise a warning. To silence it use `x.astype(y.dtype)`.
Deprecate `jax.xla_computation`. Use JAX AOT APIs to get the equivalent of `jax.xla_computation` functionality. PiperOrigin-RevId: 644107276
2024-06-17 13:01:37 -07:00
* `jax.xla_computation` is deprecated and will be removed in a future 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).
* For cross-backend lowering, you can replace
`jax.xla_computation(fn, backend='tpu')(*args, **kwargs)` with
`jax.jit(fn).trace(*args, **kwargs).lower(lowering_platforms=('tpu',)).compiler_ir('hlo')`.
[Rollback] Bumped the minimum ml_dtypes version to 0.4.0 Reverts e86c436e7f8e4e0546eff8bc2d3756a7c49dc83b PiperOrigin-RevId: 642741832
2024-06-12 14:39:11 -07:00
Finish jax and jaxlib 0.4.30 release PiperOrigin-RevId: 644402635
2024-06-18 08:55:26 -07:00
## jaxlib 0.4.30 (June 18, 2024)
Update version numbers after v0.4.29 release.
2024-06-10 14:33:13 -04:00
Update JAX dependencies, extras, and documentation for plugins. * Make jaxlib a direct dependency of jax. * Remove mentions of monolithic CUDA installations from the JAX documentation. * Drop the cuda12_pip extra and the cudnn version specific extras. * Add a with_cuda extra to the jax-cuda12-plugin package, use it in jax's setup.py. This allows us to specify cuda extras in one place. * Make a few small doc improvements.
2024-06-13 10:50:24 -04:00
* Support for monolithic CUDA jaxlibs has been dropped. You must use the
plugin-based installation (`pip install jax[cuda12]` or
`pip install jax[cuda12_local]`).
Update version numbers after v0.4.29 release.
2024-06-10 14:33:13 -04:00
## jax 0.4.29 (June 10, 2024)
Finish jax and jaxlib 0.4.28 release PiperOrigin-RevId: 632653310
2024-05-10 18:06:08 -07:00
Update release notes for 0.4.29 release.
2024-06-06 11:13:14 -04:00
* Changes
* 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]`).
Bumped the minimum ml_dtypes version to 0.4.0
2024-05-23 20:42:21 +01:00
* JAX now requires ml_dtypes version 0.4.0 or newer.
[export] Remove old deprecated APIs for jax.experimental.export. See CHANGELOG.md. The deprecation period has passed. Also replace deprecated .call_exported with .call in tests. PiperOrigin-RevId: 641236222
2024-06-07 06:51:26 -07:00
* 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.
tree_all: add support for is_leaf
2024-06-10 09:46:15 -07:00
* Added `is_leaf` argument to {func}`jax.tree.all` & {func}`jax.tree_util.tree_all`.
Bumped the minimum ml_dtypes version to 0.4.0
2024-05-23 20:42:21 +01:00
Finalize deprecations of several APIs PiperOrigin-RevId: 633634215
2024-05-14 10:39:43 -07:00
* Deprecations
Deprecate `XLACompatibleSharding` in favor of `jax.sharding.Sharding`. PiperOrigin-RevId: 640544939
2024-06-05 09:06:36 -07:00
* `jax.sharding.XLACompatibleSharding` is deprecated. Please use
`jax.sharding.Sharding`.
[export] Rename in_shardings and out_shardings fields. We rename `in_shardings` to `in_shardings_hlo` to remove confusion with JAX's use of `in_shardings`. We also rename `xla_compatible_in_sharding` to `in_shardings_jax` since we do not have a XLACompatibleSharding type anymore.
2024-06-06 18:43:25 +01:00
* `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.
Finalize deprecations of several APIs PiperOrigin-RevId: 633634215
2024-05-14 10:39:43 -07:00
* 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`.
Rename rcond/tol to rtol in linalg.matrix_rank and linalg.pinv
2024-05-14 19:53:54 +00:00
* 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.
Finalize deprecation of the config module. To configure JAX, use `import jax` and reference the config object via `jax.config`. PiperOrigin-RevId: 635430169
2024-05-20 05:48:49 -07:00
* The deprecated `jax.config` submodule has been removed. To configure JAX
use `import jax` and then reference the config object via `jax.config`.
Finalize deprecation of batched keys to PRNG functions PiperOrigin-RevId: 636196573
2024-05-22 09:36:43 -07:00
* {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.
jax.scipy.special.beta: deprecate x,y in favor of a,b
2024-06-10 09:01:39 -07:00
* In {func}`jax.scipy.special.beta`, the `x` and `y` parameters have been
renamed to `a` and `b` for consistency with other `beta` APIs.
Async dispatch expensive computations on the JAX CPU backend. By setting `jax.config.update('jax_cpu_enable_async_dispatch', False)`, one could opt out of the change and recover the old behavior. PiperOrigin-RevId: 633264117
2024-05-13 10:52:31 -07:00
[export] Rename in_shardings and out_shardings fields. We rename `in_shardings` to `in_shardings_hlo` to remove confusion with JAX's use of `in_shardings`. We also rename `xla_compatible_in_sharding` to `in_shardings_jax` since we do not have a XLACompatibleSharding type anymore.
2024-06-06 18:43:25 +01:00
* 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.
Update version numbers after v0.4.29 release.
2024-06-10 14:33:13 -04:00
## jaxlib 0.4.29 (June 10, 2024)
Finish jax and jaxlib 0.4.28 release PiperOrigin-RevId: 632653310
2024-05-10 18:06:08 -07:00
Add note to release notes about #21403. Fixes #21403
2024-05-24 10:09:13 -04:00
* Bug fixes
Update release notes for 0.4.29 release.
2024-06-06 11:13:14 -04:00
* Fixed a bug where XLA sharded some concatenation operations incorrectly,
Add note to release notes about #21403. Fixes #21403
2024-05-24 10:09:13 -04:00
which manifested as an incorrect output for cumulative reductions (#21403).
Update release notes for 0.4.29 release.
2024-06-06 11:13:14 -04:00
* Fixed a bug where XLA:CPU miscompiled certain matmul fusions
(https://github.com/openxla/xla/pull/13301).
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* Fixes a compiler crash on GPU (https://github.com/jax-ml/jax/issues/21396).
Add note to release notes about #21403. Fixes #21403
2024-05-24 10:09:13 -04:00
Issue a warning where code relies on a bug where treedef.flatten_up_to(...) was overly permissive for None treedefs. For example, tree_map(..., None, [2, 3]) previously did not raise an error, but None is a container and only leaves can be considered tree prefixes in this case. In a future release of JAX, this behavior will become an error. PiperOrigin-RevId: 641690427
2024-06-09 09:17:42 -07:00
* Deprecations
* `jax.tree.map(f, None, non-None)` now emits a `DeprecationWarning`, and will
raise an error in a future version of jax. `None` is only a tree-prefix of
itself. To preserve the current behavior, you can ask `jax.tree.map` to
treat `None` as a leaf value by writing:
`jax.tree.map(lambda x, y: None if x is None else f(x, y), a, b, is_leaf=lambda x: x is None)`.
Finish jax and jaxlib 0.4.28 release PiperOrigin-RevId: 632653310
2024-05-10 18:06:08 -07:00
## jax 0.4.28 (May 9, 2024)
Finish jax and jaxlib 0.4.27 release PiperOrigin-RevId: 631486157
2024-05-07 11:13:21 -07:00
Prepare 0.4.28 release.
2024-05-08 18:55:53 +00:00
* Bug fixes
* Reverted a change to `make_jaxpr` that was breaking Equinox (#21116).
CHANGELOG: move change from 0.4.27 to 0.4.28
2024-05-07 11:16:11 -07:00
* Deprecations & removals
* The ``kind`` argument to {func}`jax.numpy.sort` and {func}`jax.numpy.argsort`
is now removed. Use `stable=True` or `stable=False` instead.
Removed get_compute_capability from jax.experimental.pallas.gpu Compute capability is available as a `str` attribute on a GPU device since jaxlib 0.4.26.
2024-05-08 20:29:18 +01:00
* Removed ``get_compute_capability`` from the ``jax.experimental.pallas.gpu``
module. Use the ``compute_capability`` attribute of a GPU device, returned
by {func}`jax.devices` or {func}`jax.local_devices`, instead.
Deprecate newshape argument of jnp.reshape
2024-05-09 21:02:07 +00:00
* The ``newshape`` argument to {func}`jax.numpy.reshape`is being deprecated
and will soon be removed. Use `shape` instead.
CHANGELOG: move change from 0.4.27 to 0.4.28
2024-05-07 11:16:11 -07:00
Prepare 0.4.28 release.
2024-05-08 18:55:53 +00:00
* Changes
* The minimum jaxlib version of this release is 0.4.27.
Finish jax and jaxlib 0.4.28 release PiperOrigin-RevId: 632653310
2024-05-10 18:06:08 -07:00
## jaxlib 0.4.28 (May 9, 2024)
Prepare 0.4.28 release.
2024-05-08 18:55:53 +00:00
* Bug fixes
[XLA:Python] Fix a memory corruption bug in the tp_name attribute of ArrayImpl and PjitFunction for Python 3.10 or earlier. This works around https://github.com/python/cpython/issues/89478, which was fixed in Python 3.11. PiperOrigin-RevId: 631984256
2024-05-08 18:04:40 -07:00
* Fixes a memory corruption bug in the type name of Array and JIT Python
objects in Python 3.10 or earlier.
Prepare 0.4.28 release.
2024-05-08 18:55:53 +00:00
* Fixed a warning `'+ptx84' is not a recognized feature for this target`
under CUDA 12.4.
* Fixed a slow compilation problem on CPU.
* Changes
* The Windows build is now built with Clang instead of MSVC.
Finish jax and jaxlib 0.4.27 release PiperOrigin-RevId: 631486157
2024-05-07 11:13:21 -07:00
## jax 0.4.27 (May 7, 2024)
Finish jax and jaxlib 0.4.26 release PiperOrigin-RevId: 621658207
2024-04-03 15:33:20 -07:00
Add new unstack function to numpy/array_api namespaces
2024-04-15 21:03:26 +00:00
* New Functionality
Add new cumulative_sum function to numpy and array_api
2024-04-16 19:57:55 +00:00
* Added {func}`jax.numpy.unstack` and {func}`jax.numpy.cumulative_sum`,
following their addition in the array API 2023 standard, soon to be
adopted by NumPy.
Add a common flag for the collectives implementations on cpu.
2024-04-19 17:51:17 +02:00
* Added a new config option `jax_cpu_collectives_implementation` to select the
implementation of cross-process collective operations used by the CPU backend.
Choices available are `'none'`(default), `'gloo'` and `'mpi'` (requires jaxlib 0.4.26).
If set to `'none'`, cross-process collective operations are disabled.
Add new unstack function to numpy/array_api namespaces
2024-04-15 21:03:26 +00:00
jax.pure_callback and jax.experimental.io_callback now use jax.Arrays The motivation for this change is two-fold * JAX APIs should use jax.Arrays. * Using jax.Arrays potentially allows keeping the data on device, instead of always copying it to the host. Note that the version here still always copies to the host. If this change breaks you, you can recover the old behavior by changing jax.pure_callback( f, result_shape_dtypes, *args, **kwargs, ) to jax.pure_callback( lambda *args: f(*jax.tree.map(np.asarray, args)), result_shape_dtypes, *args, **kwargs, ) so that the callback function is called with NumPy arrays as before. I will update the "External callbacks" tutorial in a follow up. PiperOrigin-RevId: 622457378
2024-04-06 09:29:16 -07:00
* Changes
jax.debug.callback now passes arguments as jax.Arrays Prior to this change the behavior in eager and under jax.jit was inconsistent >>> (lambda *args: jax.debug.callback(print, *args))([42]) [42] >>> jax.jit(lambda *args: jax.debug.callback(print, *args))([42]) [array(42, dtype=int32)] It was also inconsistent with other callback APIs, which cast the arguments to jax.Arrays. Closes #20627. PiperOrigin-RevId: 626461904
2024-04-19 13:56:28 -07:00
* {func}`jax.pure_callback`, {func}`jax.experimental.io_callback`
and {func}`jax.debug.callback` now use {class}`jax.Array` instead
of {class}`np.ndarray`. You can recover the old behavior by transforming
the arguments via `jax.tree.map(np.asarray, args)` before passing them
to the callback.
Make complex_arr.astype(bool) follow NumPy's semantics
2024-04-09 16:15:59 -07:00
* `complex_arr.astype(bool)` now follows the same semantics as NumPy, returning
False where `complex_arr` is equal to `0 + 0j`, and True otherwise.
Make `core.Token` a non-trivial class which wraps a `jax.Array`. Currently, we use a singleton and empty `core.token` object everywhere. After the change, tokens could be created and threaded in and out of computations to build up dependency. Also update ordered side-effects to use the new `core.Token` class (NFC for this part, just to unify token usage). PiperOrigin-RevId: 626091210
2024-04-18 11:09:02 -07:00
* `core.Token` now is a non-trivial class which wraps a `jax.Array`. It could
be created and threaded in and out of computations to build up dependency.
The singleton object `core.token` has been removed, users now should create
and use fresh `core.Token` objects instead.
reintroduce the Threefry GPU kernel lowering, under a flag On GPU, the Threefry PRNG implementation no longer lowers to a kernel call by default. This choice can improve runtime memory usage at a compile-time cost. Prior behavior, which produces a kernel call, can be recovered with: `jax.config.update('jax_threefry_gpu_kernel_lowering', True)` PiperOrigin-RevId: 629763763
2024-05-01 10:32:36 -07:00
* On GPU, the Threefry PRNG implementation no longer lowers to a kernel call
by default. This choice can improve runtime memory usage at a compile-time
cost. Prior behavior, which produces a kernel call, can be recovered with
`jax.config.update('jax_threefry_gpu_kernel_lowering', True)`. If the new
default causes issues, please file a bug. Otherwise, we intend to remove
this flag in a future release.
jax.pure_callback and jax.experimental.io_callback now use jax.Arrays The motivation for this change is two-fold * JAX APIs should use jax.Arrays. * Using jax.Arrays potentially allows keeping the data on device, instead of always copying it to the host. Note that the version here still always copies to the host. If this change breaks you, you can recover the old behavior by changing jax.pure_callback( f, result_shape_dtypes, *args, **kwargs, ) to jax.pure_callback( lambda *args: f(*jax.tree.map(np.asarray, args)), result_shape_dtypes, *args, **kwargs, ) so that the callback function is called with NumPy arrays as before. I will update the "External callbacks" tutorial in a follow up. PiperOrigin-RevId: 622457378
2024-04-06 09:29:16 -07:00
Pallas now exclusively uses XLA for compiling kernels on GPU The old lowering pass via Triton Python APIs has been removed and the JAX_TRITON_COMPILE_VIA_XLA environment variable no longer has any effect. PiperOrigin-RevId: 621857046
2024-04-04 07:46:34 -07:00
* Deprecations & Removals
* Pallas now exclusively uses XLA for compiling kernels on GPU. The old
lowering pass via Triton Python APIs has been removed and the
`JAX_TRITON_COMPILE_VIA_XLA` environment variable no longer has any effect.
Update `jnp.clip` to Array API 2023 standard
2024-04-04 22:55:10 +00:00
* {func}`jax.numpy.clip` has a new argument signature: `a`, `a_min`, and
Fix doc Typos
2024-04-22 10:32:51 +05:30
`a_max` are deprecated in favor of `x` (positional only), `min`, and
Update `jnp.clip` to Array API 2023 standard
2024-04-04 22:55:10 +00:00
`max` ({jax-issue}`20550`).
Finalize the deprecation of the arr.device() method The method has been emitting an DeprecationWarning since JAX v0.4.21, released December 2023. Existing uses can be replaced with `arr.devices()` or `arr.sharding`, depending on the context. PiperOrigin-RevId: 623015500
2024-04-08 19:03:17 -07:00
* The `device()` method of JAX arrays has been removed, after being deprecated
since JAX v0.4.21. Use `arr.devices()` instead.
softmax: deprecate initial argument & always set to -inf internally
2024-04-10 10:23:21 -07:00
* The `initial` argument to {func}`jax.nn.softmax` and {func}`jax.nn.log_softmax`
is deprecated; empty inputs to softmax are now supported without setting this.
Finalize deprecation of invalid JIT argument names & numbers Invalid static_argnames/static_argnums have been resulting in a warning since JAX v0.3.17, released in June 2022. After this change, they will result in an error. PiperOrigin-RevId: 624270701
2024-04-12 13:08:31 -07:00
* In {func}`jax.jit`, passing invalid `static_argnums` or `static_argnames`
now leads to an error rather than a warning.
Bumped the minimum jaxlib to 0.4.23 jaxlib 0.4.23 has xla_extension_version 223 and mlir_api_version 54.
2024-04-04 14:50:42 +01:00
* The minimum jaxlib version is now 0.4.23.
Fixed hypot bug on nan/inf pairings, began deprecation of non-real values
2024-04-15 17:56:16 +00:00
* The {func}`jax.numpy.hypot` function now issues a deprecation warning when
passing complex-valued inputs to it. This will raise an error when the
deprecation is completed.
Finalize deprecation of zero-dimensional inputs to jnp.nonzero PiperOrigin-RevId: 626299531
2024-04-19 02:18:16 -07:00
* Scalar arguments to {func}`jax.numpy.nonzero`, {func}`jax.numpy.where`, and
related functions now raise an error, following a similar change in NumPy.
Add a common flag for the collectives implementations on cpu.
2024-04-19 17:51:17 +02:00
* The config option `jax_cpu_enable_gloo_collectives` is deprecated.
Use `jax.config.update('jax_cpu_collectives_implementation', 'gloo')` instead.
Finalize deprecation of `arr.device_buffer` and `arr.device_buffers` PiperOrigin-RevId: 627899901
2024-04-24 17:26:38 -07:00
* The `jax.Array.device_buffer` and `jax.Array.device_buffers` methods have
been removed after being deprecated in JAX v0.4.22. Instead use
{attr}`jax.Array.addressable_shards` and {meth}`jax.Array.addressable_data`.
Finalize deprecation of jnp.where keyword arguments PiperOrigin-RevId: 629086639
2024-04-29 09:09:13 -07:00
* The `condition`, `x`, and `y` parameters of `jax.numpy.where` are now
positional-only, following deprecation of the keywords in JAX v0.4.21.
Finalize deprecation of lax.linalg positional args PiperOrigin-RevId: 629581163
2024-04-30 17:55:31 -07:00
* Non-array arguments to functions in {mod}`jax.lax.linalg` now must be
specified by keyword. Previously, this raised a DeprecationWarning.
Require arraylike input for several jax.numpy functions PiperOrigin-RevId: 630532821
2024-05-03 16:54:22 -07:00
* Array-like arguments are now required in several :func:`jax.numpy` APIs,
including {func}`~jax.numpy.apply_along_axis`,
{func}`~jax.numpy.apply_over_axes`, {func}`~jax.numpy.inner`,
{func}`~jax.numpy.outer`, {func}`~jax.numpy.cross`,
{func}`~jax.numpy.kron`, and {func}`~jax.numpy.lexsort`.
Pallas now exclusively uses XLA for compiling kernels on GPU The old lowering pass via Triton Python APIs has been removed and the JAX_TRITON_COMPILE_VIA_XLA environment variable no longer has any effect. PiperOrigin-RevId: 621857046
2024-04-04 07:46:34 -07:00
Add support for copy kwarg in astype to match Array API
2024-04-22 16:25:37 +00:00
* Bug fixes
* {func}`jax.numpy.astype` will now always return a copy when `copy=True`.
Previously, no copy would be made when the output array would have the same
dtype as the input array. This may result in some increased memory usage.
Finalize deprecations of several APIs PiperOrigin-RevId: 633634215
2024-05-14 10:39:43 -07:00
The default value is set to `copy=False` to preserve backwards compatibility.
Add support for copy kwarg in astype to match Array API
2024-04-22 16:25:37 +00:00
Finish jax and jaxlib 0.4.27 release PiperOrigin-RevId: 631486157
2024-05-07 11:13:21 -07:00
## jaxlib 0.4.27 (May 7, 2024)
Finish jax and jaxlib 0.4.26 release PiperOrigin-RevId: 621658207
2024-04-03 15:33:20 -07:00
## jax 0.4.26 (April 3, 2024)
Finish jax and jaxlib 0.4.25 release PiperOrigin-RevId: 610413312
2024-02-26 08:16:41 -08:00
Add jax.numpy.trapezoid This function has been added to NumPy in version 2.0, as a replacement for the already deprecated trapz function.
2024-04-01 13:05:20 -07:00
* New Functionality
* Added {func}`jax.numpy.trapezoid`, following the addition of this function in
NumPy 2.0.
jnp.geomspace: make complex behavior consistent with NumPy 2.0
2024-04-02 16:12:49 -07:00
* Changes
* Complex-valued {func}`jax.numpy.geomspace` now chooses the logarithmic spiral
branch consistent with that of NumPy 2.0.
changelog: batching rule change for `rng_bit_generator`
2024-04-04 16:34:10 -07:00
* The behavior of `lax.rng_bit_generator`, and in turn the `'rbg'`
and `'unsafe_rbg'` PRNG implementations, under `jax.vmap` [has
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
changed](https://github.com/jax-ml/jax/issues/19085) so that
changelog: batching rule change for `rng_bit_generator`
2024-04-04 16:34:10 -07:00
mapping over keys results in random generation only from the first
key in the batch.
changelog: note doc change to use `jax.random.key` over `PRNGKey`
2024-04-04 16:38:08 -07:00
* Docs now use `jax.random.key` for construction of PRNG key arrays
rather than `jax.random.PRNGKey`.
jnp.geomspace: make complex behavior consistent with NumPy 2.0
2024-04-02 16:12:49 -07:00
Remove --jax_parallel_functions_output_gda. PiperOrigin-RevId: 616898032
2024-03-18 11:41:17 -07:00
* Deprecations & Removals
Deprecate jax.tree_map for jax v0.4.26 Reverts f4045dceb206be1ea10ee651ccc6151809f2d9f3 PiperOrigin-RevId: 611230367
2024-02-28 14:28:17 -08:00
* {func}`jax.tree_map` is deprecated; use `jax.tree.map` instead, or for backward
compatibility with older JAX versions, use {func}`jax.tree_util.tree_map`.
Deprecate `jax.clear_backends`. `jax.clear_backends` does not necessarily do what its name suggests and can lead to unexpected consequences, e.g., it will not destroy existing backends and release corresponding owned resources. Use `jax.clear_caches` if you only want to clean up compilation caches. For backward compatibilty or you really need to switch/reinitialize the default backend, use `jax.extend.backend.clear_backends`. PiperOrigin-RevId: 616946337
2024-03-18 14:22:35 -07:00
* {func}`jax.clear_backends` is deprecated as it does not necessarily do what
its name suggests and can lead to unexpected consequences, e.g., it will not
destroy existing backends and release corresponding owned resources. Use
{func}`jax.clear_caches` if you only want to clean up compilation caches.
For backward compatibility or you really need to switch/reinitialize the
default backend, use {func}`jax.extend.backend.clear_backends`.
Deprecated the jax.experimental.maps submodule PiperOrigin-RevId: 614082251
2024-03-08 16:50:02 -08:00
* The `jax.experimental.maps` module and `jax.experimental.maps.xmap` are
deprecated. Use `jax.experimental.shard_map` or `jax.vmap` with the
`spmd_axis_name` argument for expressing SPMD device-parallel computations.
[host_callback] Deprecate the jax.experimental.host_callback module.
2024-03-21 08:18:57 +02:00
* The `jax.experimental.host_callback` module is deprecated.
Use instead the [new JAX external callbacks](https://jax.readthedocs.io/en/latest/notebooks/external_callbacks.html).
[callback] Add a flag to implement host_callback in terms of io_callback. The host_callbacks APIs are deprecated and will be removed. In order to help the transition to the new APIs, we add a flag (`JAX_HOST_CALLBACK_LEGACY`) that when set to `False` will use `io_callback` (and `pure_callback` and `jax.debug.callback`) to implement the host_callback APIs. See issue #20385 for more details. We change the tests to accomodate slightly different results when using the new callbacks. The tests that use `tap_with_device` and `call_with_device` are disabled when using the new callbacks.
2024-03-29 13:36:20 +02:00
Added `JAX_HOST_CALLBACK_LEGACY` flag to assist in the transition to the
new callbacks. See {jax-issue}`#20385` for a discussion.
Finalize deprecation of non-array arguments to array_equal/array_equiv
2024-02-27 13:48:08 -08:00
* Passing arguments to {func}`jax.numpy.array_equal` and {func}`jax.numpy.array_equiv`
that cannot be converted to a JAX array now results in an exception.
Remove --jax_parallel_functions_output_gda. PiperOrigin-RevId: 616898032
2024-03-18 11:41:17 -07:00
* The deprecated flag `jax_parallel_functions_output_gda` has been removed.
This flag was long deprecated and did nothing; its use was a no-op.
Finalize deprecations of jax.interpreters.ad config & source_info_util These have been raising a DeprecationWarning since JAX 0.4.19, released 2023 Oct 19. I've left the undefined symbols in place for now, as they will raise an informative AttributeError. PiperOrigin-RevId: 616931120
2024-03-18 13:32:10 -07:00
* The previously-deprecated imports `jax.interpreters.ad.config` and
`jax.interpreters.ad.source_info_util` have now been removed. Use `jax.config`
and `jax.extend.source_info_util` instead.
Fix doc typos
2024-04-05 14:21:33 +05:30
* JAX export does not support older serialization versions anymore. Version 9
[export] Increase minimum serialization version to 9. Stop supporting serializing older version. The current max serialization version 9 has been supported since October 27th, 2023 and has become the default since February 1, 2024. This change could break clients that set a specific JAX serialization version lower than 9. See https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#native-serialization-versions PiperOrigin-RevId: 619588685
2024-03-27 11:05:25 -07:00
has been supported since October 27th, 2023 and has become the default
since February 1, 2024.
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
See [a description of the versions](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/README.md#native-serialization-versions).
[export] Increase minimum serialization version to 9. Stop supporting serializing older version. The current max serialization version 9 has been supported since October 27th, 2023 and has become the default since February 1, 2024. This change could break clients that set a specific JAX serialization version lower than 9. See https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#native-serialization-versions PiperOrigin-RevId: 619588685
2024-03-27 11:05:25 -07:00
This change could break clients that set a specific
JAX serialization version lower than 9.
Finish jax and jaxlib 0.4.26 release PiperOrigin-RevId: 621658207
2024-04-03 15:33:20 -07:00
## jaxlib 0.4.26 (April 3, 2024)
Finish jax and jaxlib 0.4.25 release PiperOrigin-RevId: 610413312
2024-02-26 08:16:41 -08:00
Remove support for CUDA 11. Pin minimal required versions for CUDA to 12.1. Reverts 910a31d7b7510e3375718ab1ea0d38df7bd2c0d5 PiperOrigin-RevId: 618911489
2024-03-25 11:44:40 -07:00
* Changes
* JAX now supports CUDA 12.1 or newer only. Support for CUDA 11.8 has been
dropped.
Prepare for 0.4.26 release.
2024-04-03 14:38:58 -04:00
* JAX now supports NumPy 2.0.
Remove support for CUDA 11. Pin minimal required versions for CUDA to 12.1. Reverts 910a31d7b7510e3375718ab1ea0d38df7bd2c0d5 PiperOrigin-RevId: 618911489
2024-03-25 11:44:40 -07:00
Finish jax and jaxlib 0.4.25 release PiperOrigin-RevId: 610413312
2024-02-26 08:16:41 -08:00
## jax 0.4.25 (Feb 26, 2024)
Update version numbers post-0.4.24 release
2024-02-06 16:28:23 -08:00
Add CUDA Array Interface consumer support
2024-01-07 19:15:32 +02:00
* New Features
* Added [CUDA Array
Interface](https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html)
Fix jax 0.4.24 changelog
2024-02-08 14:17:31 -08:00
import support (requires jaxlib 0.4.24).
jax.numpy: implement scalar boolean indexing
2024-02-08 14:11:15 -08:00
* JAX arrays now support NumPy-style scalar boolean indexing, e.g. `x[True]` or `x[False]`.
Add jax.tree module with aliases of jax.tree_util
2024-02-12 13:07:59 -08:00
* Added {mod}`jax.tree` module, with a more convenient interface for referencing functions
in {mod}`jax.tree_util`.
tree_transpose: optionally infer inner_treedef
2024-02-15 12:01:21 -08:00
* {func}`jax.tree.transpose` (i.e. {func}`jax.tree_util.tree_transpose`) now accepts
`inner_treedef=None`, in which case the inner treedef will be automatically inferred.
Fix jax 0.4.24 changelog
2024-02-08 14:17:31 -08:00
Compile Triton kernels via XLA by default PiperOrigin-RevId: 609299269
2024-02-22 02:31:33 -08:00
* Changes
* Pallas now uses XLA instead of the Triton Python APIs to compile Triton
kernels. You can revert to the old behavior by setting the
`JAX_TRITON_COMPILE_VIA_XLA` environment variable to `"0"`.
Mention re-instated xla APIs in the CHANGELOG
2024-02-22 12:19:29 -08:00
* Several deprecated APIs in {mod}`jax.interpreters.xla` that were removed in v0.4.24
have been re-added in v0.4.25, including `backend_specific_translations`,
`translations`, `register_translation`, `xla_destructure`, `TranslationRule`,
`TranslationContext`, and `XLAOp`. These are still considered deprecated, and
will be removed again in the future when better replacements are available.
Refer to {jax-issue}`#19816` for discussion.
Compile Triton kernels via XLA by default PiperOrigin-RevId: 609299269
2024-02-22 02:31:33 -08:00
Fix jax 0.4.24 changelog
2024-02-08 14:17:31 -08:00
* Deprecations & Removals
* {func}`jax.numpy.linalg.solve` now shows a deprecation warning for batched 1D
solves with `b.ndim > 1`. In the future these will be treated as batched 2D
solves.
Error for deprecated scalar conversions of non-scalar arrays
2024-02-15 16:55:33 -08:00
* Conversion of a non-scalar array to a Python scalar now raises an error, regardless
of the size of the array. Previously a deprecation warning was raised in the case of
non-scalar arrays of size 1. This follows a similar deprecation in NumPy.
Removed deprecated jax.config methods and jax.config.config Reverts dcc65e621ea3a68fdc79fa9f2c995743a7b3faf7 PiperOrigin-RevId: 608676645
2024-02-20 11:24:40 -08:00
* The previously deprecated configuration APIs have been removed
following a standard 3 months deprecation cycle (see {ref}`api-compatibility`).
These include
* the `jax.config.config` object and
* the `define_*_state` and `DEFINE_*` methods of {data}`jax.config`.
* Importing the `jax.config` submodule via `import jax.config` is deprecated.
To configure JAX use `import jax` and then reference the config object
via `jax.config`.
Increase minimum jaxlib version to 0.4.20. jaxlib 0.4.20 has xla_extension_version 210 and mlir_api_version 54. PiperOrigin-RevId: 609094229
2024-02-21 12:58:18 -08:00
* The minimum jaxlib version is now 0.4.20.
Fix jax 0.4.24 changelog
2024-02-08 14:17:31 -08:00
Finish jax and jaxlib 0.4.25 release PiperOrigin-RevId: 610413312
2024-02-26 08:16:41 -08:00
## jaxlib 0.4.25 (Feb 26, 2024)
Fix jax 0.4.24 changelog
2024-02-08 14:17:31 -08:00
## jax 0.4.24 (Feb 6, 2024)
Add CUDA Array Interface consumer support
2024-01-07 19:15:32 +02:00
[roll forward 2] Remove the `jax_require_devices_during_lowering flag since it was temporary. Added the semi-breaking change to Changelog.md. Reverts b52bcc1639368069075284eefc763f824ca155f1 PiperOrigin-RevId: 590959383
2023-12-14 09:13:43 -08:00
* Changes
[export] Refactor the imports for the public API of jax.experimental.export Previously we used `from jax.experimental.export import export` and `export.export(fun)`. Now we want to add the public API directly to `jax.experimental.export`, for the following desired usage: ``` from jax.experimental import export exp: export.Exported = export.export(fun) ser: bytearray = export.serialize(exp) exp1: export.Exported = export.deserialized(ser) export.call(exp1) ``` This change requires changing the type of `jax.experimental.export.export` from a module to a function. This confuses pytype for the targets with strict type checking, which is why I attempt to make this change atomically throughout the internal code base. In order to support backwards compatibility with OSS packages, this change also includes explicit JAX version checks in several OSS packages, and also adds to the `export` function the attributes that the old export module had. PiperOrigin-RevId: 596563481
2024-01-08 05:29:11 -08:00
[roll forward 2] Remove the `jax_require_devices_during_lowering flag since it was temporary. Added the semi-breaking change to Changelog.md. Reverts b52bcc1639368069075284eefc763f824ca155f1 PiperOrigin-RevId: 590959383
2023-12-14 09:13:43 -08:00
* JAX lowering to StableHLO does not depend on physical devices anymore.
Finalize deprecations of several APIs PiperOrigin-RevId: 633634215
2024-05-14 10:39:43 -07:00
If your primitive wraps custom_partitioning or JAX callbacks in the lowering
[roll forward 2] Remove the `jax_require_devices_during_lowering flag since it was temporary. Added the semi-breaking change to Changelog.md. Reverts b52bcc1639368069075284eefc763f824ca155f1 PiperOrigin-RevId: 590959383
2023-12-14 09:13:43 -08:00
rule i.e. function passed to `rule` parameter of `mlir.register_lowering` then add your
primitive to `jax._src.dispatch.prim_requires_devices_during_lowering` set.
This is needed because custom_partitioning and JAX callbacks need physical
devices to create `Sharding`s during lowering.
This is a temporary state until we can create `Sharding`s without physical
devices.
[array api] add stable & descending params to jnp.sort & jnp.argsort
2024-01-04 14:21:25 -08:00
* {func}`jax.numpy.argsort` and {func}`jax.numpy.sort` now support the `stable`
and `descending` arguments.
[shape_poly] Replace non_negative_dim with max_dim and min_dim. Previously, we had `core.non_negative_dim` and we used it to express `max(d, 0)`. This is needed in several places internally to express index computations involving clamping (for numpy indexing), or striding and dilation (which have a conditional semantics). It seemed that this special case was sufficient, and we expressed `max(a, b)` as `a + non_negative(b - a)` and `min(a, b)` as `a - non_negative(a - b)`. One drawback was that `non_negative` can be a surprising construct when it appears in error messages. Also, users need `max` and `min` computations with dimensions. It is clearer if we use `max` and `min` directly instead of rewriting these to use `non_negative`. The drawback is that we now have to duplicate some internal logic to for `max` and `min`, but overall I feel this is worth it for the better error messages we get.
2023-12-13 10:14:27 +01:00
* Several changes to the handling of shape polymorphism (used in
[shape_poly] Introduce is_symbolic_dim and deprecate is_poly_dim. The old is_poly_dim seems to be used in a few places externally. This was from the time when the symbolic dimensions were polynomials, now we use the more generic term symbolic dimension or expression. We introduce is_symbolic_dim and we export it through the jax.experimental.export. We plan to make the entire shape_poly.py module private, and this is a necessary step.
2024-01-10 08:45:03 +02:00
{mod}`jax.experimental.jax2tf` and {mod}`jax.experimental.export`):
* cleaner pretty-printing of symbolic expressions ({jax-issue}`#19227`)
[shape_poly] Forgot to update CHANGELOG for #19235.
2024-01-23 17:31:23 +01:00
* added the ability to specify symbolic constraints on the dimension variables.
This makes shape polymorphism more expressive, and gives a way to workaround
limitations in the reasoning about inequalities.
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
See https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/README.md#user-specified-symbolic-constraints.
[shape_poly] Forgot to update CHANGELOG for #19235.
2024-01-23 17:31:23 +01:00
* with the addition of symbolic constraints ({jax-issue}`#19235`) we now
consider dimension variables from different scopes to be different, even
if they have the same name. Symbolic expressions from different scopes
cannot interact, e.g., in arithmetic operations.
Scopes are introduced by {func}`jax.experimental.jax2tf.convert`,
{func}`jax.experimental.export.symbolic_shape`, {func}`jax.experimental.export.symbolic_args_specs`.
Bumped the minimum jaxlib to 0.4.23 jaxlib 0.4.23 has xla_extension_version 223 and mlir_api_version 54.
2024-04-04 14:50:42 +01:00
The scope of a symbolic expression `e` can be read with `e.scope` and passed
Fix doc typos
2024-04-05 14:21:33 +05:30
into the above functions to direct them to construct symbolic expressions in
[shape_poly] Forgot to update CHANGELOG for #19235.
2024-01-23 17:31:23 +01:00
a given scope.
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
See https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/README.md#user-specified-symbolic-constraints.
[shape_poly] Introduce is_symbolic_dim and deprecate is_poly_dim. The old is_poly_dim seems to be used in a few places externally. This was from the time when the symbolic dimensions were polynomials, now we use the more generic term symbolic dimension or expression. We introduce is_symbolic_dim and we export it through the jax.experimental.export. We plan to make the entire shape_poly.py module private, and this is a necessary step.
2024-01-10 08:45:03 +02:00
* simplified and faster equality comparisons, where we consider two symbolic dimensions
to be equal if the normalized form of their difference reduces to 0
({jax-issue}`#19231`; note that this may result in user-visible behavior
changes)
* improved the error messages for inconclusive inequality comparisons
[shape_poly] Deprecate shape_poly.PolySpec. This class has limited usefulness, and it seems worth removing it in favor of using strings for polymorphic specifications, thus reducing the API surface.
2024-01-10 09:05:16 +02:00
({jax-issue}`#19235`).
[shape_poly] Introduce is_symbolic_dim and deprecate is_poly_dim. The old is_poly_dim seems to be used in a few places externally. This was from the time when the symbolic dimensions were polynomials, now we use the more generic term symbolic dimension or expression. We introduce is_symbolic_dim and we export it through the jax.experimental.export. We plan to make the entire shape_poly.py module private, and this is a necessary step.
2024-01-10 08:45:03 +02:00
* the `core.non_negative_dim` API (introduced recently)
was deprecated and `core.max_dim` and `core.min_dim` were introduced
({jax-issue}`#18953`) to express `max` and `min` for symbolic dimensions.
You can use `core.max_dim(d, 0)` instead of `core.non_negative_dim(d)`.
[shape_poly] Improve and rename export.args_specs. We rename it to `symbolic_args_specs` in line with the other public APIs related to shape polymorphism. The function used to be in _export.py for historical reasons, we now move it to shape_poly.py but we export the `symbolci_args_specs` from the public `jax.experimental.export`. The improvement is that for the case when the `args` passed in are TF arrays, we move the logic to extract the shapes and dtypes from this function to the callers. This achieves a better separation of the JAX and TF use cases.
2024-01-10 09:44:31 +02:00
* the `shape_poly.is_poly_dim` is deprecated in favor of `export.is_symbolic_dim`
[shape_poly] Introduce is_symbolic_dim and deprecate is_poly_dim. The old is_poly_dim seems to be used in a few places externally. This was from the time when the symbolic dimensions were polynomials, now we use the more generic term symbolic dimension or expression. We introduce is_symbolic_dim and we export it through the jax.experimental.export. We plan to make the entire shape_poly.py module private, and this is a necessary step.
2024-01-10 08:45:03 +02:00
({jax-issue}`#19282`).
[shape_poly] Improve and rename export.args_specs. We rename it to `symbolic_args_specs` in line with the other public APIs related to shape polymorphism. The function used to be in _export.py for historical reasons, we now move it to shape_poly.py but we export the `symbolci_args_specs` from the public `jax.experimental.export`. The improvement is that for the case when the `args` passed in are TF arrays, we move the logic to extract the shapes and dtypes from this function to the callers. This achieves a better separation of the JAX and TF use cases.
2024-01-10 09:44:31 +02:00
* the `export.args_specs` is deprecated in favor of `export.symbolic_args_specs
({jax-issue}`#19283`).
[shape_poly] Deprecate shape_poly.PolySpec. This class has limited usefulness, and it seems worth removing it in favor of using strings for polymorphic specifications, thus reducing the API surface.
2024-01-10 09:05:16 +02:00
* the `shape_poly.PolyShape` and `jax2tf.PolyShape` are deprecated, use
strings for polymorphic shapes specifications ({jax-issue}`#19284`).
[export] Set default native serialization version to 9. This version adds better support for JAX effects. See description in CHANGELOG.md and also at https://github.com/google/jax/blob/main/jax/experimental/jax2tf/README.md#native-serialization-versions. PiperOrigin-RevId: 603579274
2024-02-01 21:55:31 -08:00
* JAX default native serialization version is now 9. This is relevant
for {mod}`jax.experimental.jax2tf` and {mod}`jax.experimental.export`.
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
See [description of version numbers](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/README.md#native-serialization-versions).
[export] Refactor the imports for the public API of jax.experimental.export Previously we used `from jax.experimental.export import export` and `export.export(fun)`. Now we want to add the public API directly to `jax.experimental.export`, for the following desired usage: ``` from jax.experimental import export exp: export.Exported = export.export(fun) ser: bytearray = export.serialize(exp) exp1: export.Exported = export.deserialized(ser) export.call(exp1) ``` This change requires changing the type of `jax.experimental.export.export` from a module to a function. This confuses pytype for the targets with strict type checking, which is why I attempt to make this change atomically throughout the internal code base. In order to support backwards compatibility with OSS packages, this change also includes explicit JAX version checks in several OSS packages, and also adds to the `export` function the attributes that the old export module had. PiperOrigin-RevId: 596563481
2024-01-08 05:29:11 -08:00
* Refactored the API for `jax.experimental.export`. Instead of
`from jax.experimental.export import export` you should use now
`from jax.experimental import export`. The old way of importing will
continue to work for a deprecation period of 3 months.
Add scipy.stats.sem.
2024-01-13 22:17:21 -05:00
* Added {func}`jax.scipy.stats.sem`.
jnp.unique: make return_inverse shape match NumPy 2.0
2024-01-11 12:15:52 -08:00
* {func}`jax.numpy.unique` with `return_inverse = True` returns inverse indices
reshaped to the dimension of the input, following a similar change to
{func}`numpy.unique` in NumPy 2.0.
jnp.sign: use x/abs(x) for complex arguments
2024-01-17 08:59:40 -08:00
* {func}`jax.numpy.sign` now returns `x / abs(x)` for nonzero complex inputs. This is
consistent with the behavior of {func}`numpy.sign` in NumPy version 2.0.
logsumexp: use NumPy 2.0 convention for complex sign
2024-01-16 16:15:06 -08:00
* {func}`jax.scipy.special.logsumexp` with `return_sign=True` now uses the NumPy 2.0
convention for the complex sign, `x / abs(x)`. This is consistent with the behavior
jnp.sign: use x/abs(x) for complex arguments
2024-01-17 08:59:40 -08:00
of {func}`scipy.special.logsumexp` in SciPy v1.13.
Add support for bool dlpack values. PiperOrigin-RevId: 599199196
2024-01-17 09:30:04 -08:00
* JAX now supports the bool DLPack type for both import and export.
Previously bool values could not be imported and were exported as integers.
Remove deprecated unsafe_raw_array method from PRNG keys PiperOrigin-RevId: 595190146
2024-01-02 13:02:46 -08:00
* Deprecations & Removals
Remove a number of deprecated APIs All of these were deprecated prior to the JAX 0.4.16 release, on Sept 18 2023. As of Monday Dec 18, we have met the 3 month deprecation period specified by the [API Compatiblity Policy](https://jax.readthedocs.io/en/latest/api_compatibility.html). PiperOrigin-RevId: 591933493
2023-12-18 10:08:02 -08:00
* A number of previously deprecated functions have been removed, following a
standard 3+ month deprecation cycle (see {ref}`api-compatibility`).
This includes:
Finalize a number of deprecations from JAX 0.4.19 PiperOrigin-RevId: 600509530
2024-01-22 11:00:53 -08:00
* From {mod}`jax.core`: `TracerArrayConversionError`,
`TracerIntegerConversionError`, `UnexpectedTracerError`,
`as_hashable_function`, `collections`, `dtypes`, `lu`, `map`,
`namedtuple`, `partial`, `pp`, `ref`, `safe_zip`, `safe_map`,
`source_info_util`, `total_ordering`, `traceback_util`, `tuple_delete`,
`tuple_insert`, and `zip`.
* From {mod}`jax.lax`: `dtypes`, `itertools`, `naryop`, `naryop_dtype_rule`,
`standard_abstract_eval`, `standard_naryop`, `standard_primitive`,
`standard_unop`, `unop`, and `unop_dtype_rule`.
Remove a number of deprecated APIs All of these were deprecated prior to the JAX 0.4.16 release, on Sept 18 2023. As of Monday Dec 18, we have met the 3 month deprecation period specified by the [API Compatiblity Policy](https://jax.readthedocs.io/en/latest/api_compatibility.html). PiperOrigin-RevId: 591933493
2023-12-18 10:08:02 -08:00
* The `jax.linear_util` submodule and all its contents.
* The `jax.prng` submodule and all its contents.
Finalize a number of deprecations from JAX 0.4.19 PiperOrigin-RevId: 600509530
2024-01-22 11:00:53 -08:00
* From {mod}`jax.random`: `PRNGKeyArray`, `KeyArray`, `default_prng_impl`,
`threefry_2x32`, `threefry2x32_key`, `threefry2x32_p`, `rbg_key`, and
`unsafe_rbg_key`.
Remove a number of deprecated APIs All of these were deprecated prior to the JAX 0.4.16 release, on Sept 18 2023. As of Monday Dec 18, we have met the 3 month deprecation period specified by the [API Compatiblity Policy](https://jax.readthedocs.io/en/latest/api_compatibility.html). PiperOrigin-RevId: 591933493
2023-12-18 10:08:02 -08:00
* From {mod}`jax.tree_util`: `register_keypaths`, `AttributeKeyPathEntry`, and
`GetItemKeyPathEntry`.
* from {mod}`jax.interpreters.xla`: `backend_specific_translations`, `translations`,
`register_translation`, `xla_destructure`, `TranslationRule`, `TranslationContext`,
`axis_groups`, `ShapedArray`, `ConcreteArray`, `AxisEnv`, `backend_compile`,
and `XLAOp`.
* from {mod}`jax.numpy`: `NINF`, `NZERO`, `PZERO`, `row_stack`, `issubsctype`,
Remove deprecated function jax.numpy.trapz This was deprecated prior to the JAX 0.4.16 release, so we have now met the 3 month deprecation period specified by the [API Compatiblity Policy](https://jax.readthedocs.io/en/latest/api_compatibility.html). PiperOrigin-RevId: 592266215
2023-12-19 09:57:01 -08:00
`trapz`, and `in1d`.
Remove a number of deprecated APIs All of these were deprecated prior to the JAX 0.4.16 release, on Sept 18 2023. As of Monday Dec 18, we have met the 3 month deprecation period specified by the [API Compatiblity Policy](https://jax.readthedocs.io/en/latest/api_compatibility.html). PiperOrigin-RevId: 591933493
2023-12-18 10:08:02 -08:00
* from {mod}`jax.scipy.linalg`: `tril` and `triu`.
Remove deprecated unsafe_raw_array method from PRNG keys PiperOrigin-RevId: 595190146
2024-01-02 13:02:46 -08:00
* The previously-deprecated method `PRNGKeyArray.unsafe_raw_array` has been
removed. Use {func}`jax.random.key_data` instead.
Error on conversion of empty arrays to boolean. PiperOrigin-RevId: 595264332
2024-01-02 19:26:09 -08:00
* `bool(empty_array)` now raises an error rather than returning `False`. This
previously raised a deprecation warning, and follows a similar change in NumPy.
Deprecate support for the mhlo dialect. JAX has not used mhlo for some time, in favor of stablehlo. Deprecate support for this dialect in JAX's API and remove testing. PiperOrigin-RevId: 598550225
2024-01-15 02:12:52 -08:00
* Support for the mhlo MLIR dialect has been deprecated. JAX no longer uses
the mhlo dialect, in favor of stablehlo. APIs that refer to "mhlo" will be
removed in the future. Use the "stablehlo" dialect instead.
jax.random: deprecate passing of batched keys to APIs
2024-01-17 12:53:24 -08:00
* {mod}`jax.random`: passing batched keys directly to random number generation functions,
such as {func}`~jax.random.bits`, {func}`~jax.random.gamma`, and others, is deprecated
and will emit a `FutureWarning`. Use `jax.vmap` for explicit batching.
Deprecate jax.lax.tie_in
2024-01-18 13:13:47 -08:00
* {func}`jax.lax.tie_in` is deprecated: it has been a no-op since JAX v0.2.0.
[roll forward 2] Remove the `jax_require_devices_during_lowering flag since it was temporary. Added the semi-breaking change to Changelog.md. Reverts b52bcc1639368069075284eefc763f824ca155f1 PiperOrigin-RevId: 590959383
2023-12-14 09:13:43 -08:00
Update version numbers post-0.4.24 release
2024-02-06 16:28:23 -08:00
## jaxlib 0.4.24 (Feb 6, 2024)
Finish release of jax and jaxlib 0.4.23 PiperOrigin-RevId: 590833947
2023-12-13 23:38:26 -08:00
[jax] Add pybind for `PjRtExecutable::GetCostAnalysis`/`Executable.cost_analysis` This makes cross-compiled `Compiled.cost_analysis` work. PiperOrigin-RevId: 597411014
2024-01-10 18:51:15 -08:00
* Changes
Update changelog and install doc to mention cuda 12.3 switch PiperOrigin-RevId: 605473026
2024-02-08 17:21:16 -08:00
* JAX now supports CUDA 12.3 and CUDA 11.8. Support for CUDA 12.2 has been
dropped.
[jax] Add pybind for `PjRtExecutable::GetCostAnalysis`/`Executable.cost_analysis` This makes cross-compiled `Compiled.cost_analysis` work. PiperOrigin-RevId: 597411014
2024-01-10 18:51:15 -08:00
* `cost_analysis` now works with cross-compiled `Compiled` objects (i.e. when
using `.lower().compile()` with a topology object, e.g., to compile for
Cloud TPU from a non-TPU computer).
Fix jax 0.4.24 changelog
2024-02-08 14:17:31 -08:00
* Added [CUDA Array
Interface](https://numba.readthedocs.io/en/stable/cuda/cuda_array_interface.html)
import support (requires jax 0.4.25).
[jax] Add pybind for `PjRtExecutable::GetCostAnalysis`/`Executable.cost_analysis` This makes cross-compiled `Compiled.cost_analysis` work. PiperOrigin-RevId: 597411014
2024-01-10 18:51:15 -08:00
Finish release of jax and jaxlib 0.4.23 PiperOrigin-RevId: 590833947
2023-12-13 23:38:26 -08:00
## jax 0.4.23 (Dec 13, 2023)
## jaxlib 0.4.23 (Dec 13, 2023)
Finish jax and jaxlib 0.4.22 release PiperOrigin-RevId: 590775311
2023-12-13 18:26:06 -08:00
Add patch to suppress XLA:GPU logging. PiperOrigin-RevId: 590780227
2023-12-13 18:53:01 -08:00
* Fixed a bug that caused verbose logging from the GPU compiler during
compilation.
Finish jax and jaxlib 0.4.22 release PiperOrigin-RevId: 590775311
2023-12-13 18:26:06 -08:00
## jax 0.4.22 (Dec 13, 2023)
Finish jax and jaxlib 0.4.21 release PiperOrigin-RevId: 587866580
2023-12-04 15:51:21 -08:00
Deprecate arr.device_buffer and arr.device_buffers
2023-12-06 10:20:29 -08:00
* Deprecations
* The `device_buffer` and `device_buffers` properties of JAX arrays are deprecated.
Explicit buffers have been replaced by the more flexible array sharding interface,
but the previous outputs can be recovered this way:
* `arr.device_buffer` becomes `arr.addressable_data(0)`
* `arr.device_buffers` becomes `[x.data for x in arr.addressable_shards]`
Finish jax and jaxlib 0.4.22 release PiperOrigin-RevId: 590775311
2023-12-13 18:26:06 -08:00
## jaxlib 0.4.22 (Dec 13, 2023)
Finish jax and jaxlib 0.4.21 release PiperOrigin-RevId: 587866580
2023-12-04 15:51:21 -08:00
## jax 0.4.21 (Dec 4 2023)
Finish 0.4.19 jax and jaxlib release PiperOrigin-RevId: 574983871
2023-10-19 13:27:16 -07:00
Add nn.squareplus.
2023-11-14 23:52:41 -05:00
* New Features
* Added {obj}`jax.nn.squareplus`.
Increase minimum jaxlib version to 0.4.19. 0.4.19 has xla_extension version 207 and mlir_api_version 54. PiperOrigin-RevId: 583412447
2023-11-17 09:37:45 -08:00
* Changes
* The minimum jaxlib version is now 0.4.19.
Reverts 2aaa7559f96e4bb7b0271665bf386bf3ba22c451 PiperOrigin-RevId: 584033001
2023-11-20 08:23:00 -08:00
* Released wheels are built now with clang instead of gcc.
Automate arguments for jax.distributed.initialize for cloud TPU environments. PiperOrigin-RevId: 586892544
2023-11-30 22:24:22 -08:00
* Enforce that the device backend has not been initialized prior to calling `jax.distributed.initialize()`.
* Automate arguments to `jax.distributed.initialize()` in cloud TPU environments.
Increase minimum jaxlib version to 0.4.19. 0.4.19 has xla_extension version 207 and mlir_api_version 54. PiperOrigin-RevId: 583412447
2023-11-17 09:37:45 -08:00
Remove deprecated `sym_pos` argument from `jax.scipy.linalg.solve` PiperOrigin-RevId: 580940755
2023-11-09 09:52:53 -08:00
* Deprecations
* The previously-deprecated `sym_pos` argument has been removed from
{func}`jax.scipy.linalg.solve`. Use `assume_a='pos'` instead.
Deprecate passing of None to jax.numpy.array
2023-11-16 15:10:56 -08:00
* Passing `None` to {func}`jax.array` or {func}`jax.asarray`, either directly or
within a list or tuple, is deprecated and now raises a {obj}`FutureWarning`.
It currently is converted to NaN, and in the future will raise a {obj}`TypeError`.
Make jax.numpy.where()'s condition, x, y arguments positional-only to match numpy.where. PiperOrigin-RevId: 584377134
2023-11-21 11:09:28 -08:00
* Passing the `condition`, `x`, and `y` parameters to `jax.numpy.where` by
keyword arguments has been deprecated, to match `numpy.where`.
Deprecate non-array inputs to jnp.array_equal & jnp.array_equiv
2023-11-28 13:55:18 -08:00
* Passing arguments to {func}`jax.numpy.array_equal` and {func}`jax.numpy.array_equiv`
that cannot be converted to a JAX array is deprecated and now raises a
{obj}`DeprecationWaning`. Currently the functions return False, in the future this
will raise an exception.
Fix Typos in CHANGELOG.md This PR fixes the typos in Change log documentation
2024-03-12 13:57:07 +05:30
* The `device()` method of JAX arrays is deprecated. Depending on the context, it may
Deprecate the device() method of JAX arrays
2023-11-29 16:52:09 -08:00
be replaced with one of the following:
- {meth}`jax.Array.devices` returns the set of all devices used by the array.
- {attr}`jax.Array.sharding` gives the sharding configuration used by the array.
Remove deprecated `sym_pos` argument from `jax.scipy.linalg.solve` PiperOrigin-RevId: 580940755
2023-11-09 09:52:53 -08:00
Finish jax and jaxlib 0.4.21 release PiperOrigin-RevId: 587866580
2023-12-04 15:51:21 -08:00
## jaxlib 0.4.21 (Dec 4 2023)
Update versions and CHANGELOG after jax 0.4.20 release
2023-11-02 16:30:56 -07:00
Use a Jacobi SVD solver for unbatched SVDs up to 1024x1024 on NVIDIA GPUs. The unbatched Jacobi solver is faster for small-moderate matrices, and the unbatched kernel doesn't have size restrictions. Timings on T4 GPU: Before: ------------------------------------------------------------ Benchmark Time CPU Iterations ------------------------------------------------------------ svd/m:1/n:1 263587 ns 242274 ns 2780 svd/m:2/n:1 335561 ns 298238 ns 2303 svd/m:5/n:1 337784 ns 299841 ns 2304 svd/m:10/n:1 339184 ns 300703 ns 2311 svd/m:100/n:1 359826 ns 320088 ns 2159 svd/m:500/n:1 376124 ns 338660 ns 2076 svd/m:800/n:1 375779 ns 335590 ns 2060 svd/m:1000/n:1 419171 ns 341487 ns 2072 svd/m:1/n:2 307564 ns 270663 ns 2544 svd/m:2/n:2 320928 ns 283601 ns 2487 svd/m:5/n:2 377373 ns 344228 ns 2035 svd/m:10/n:2 380557 ns 349412 ns 1953 svd/m:100/n:2 435465 ns 403496 ns 1722 svd/m:500/n:2 444610 ns 410913 ns 1680 svd/m:800/n:2 454493 ns 416495 ns 1665 svd/m:1000/n:2 492110 ns 420539 ns 1665 svd/m:1/n:5 307316 ns 275833 ns 2531 svd/m:2/n:5 374318 ns 341432 ns 2086 svd/m:5/n:5 512928 ns 470293 ns 1361 svd/m:10/n:5 589330 ns 537070 ns 1353 svd/m:100/n:5 620164 ns 580166 ns 1193 svd/m:500/n:5 636424 ns 593692 ns 1180 svd/m:800/n:5 635545 ns 595016 ns 1181 svd/m:1000/n:5 672443 ns 597387 ns 1115 svd/m:1/n:10 310013 ns 273998 ns 2520 svd/m:2/n:10 370451 ns 334489 ns 2105 svd/m:5/n:10 560037 ns 522223 ns 1274 svd/m:10/n:10 572868 ns 535388 ns 1304 svd/m:100/n:10 959802 ns 918258 ns 765 svd/m:500/n:10 955958 ns 909778 ns 758 svd/m:800/n:10 924104 ns 879512 ns 777 svd/m:1000/n:10 950140 ns 883493 ns 775 svd/m:1/n:100 351237 ns 315554 ns 2198 svd/m:2/n:100 426883 ns 390089 ns 1792 svd/m:5/n:100 601557 ns 564493 ns 1255 svd/m:10/n:100 920819 ns 880011 ns 787 svd/m:100/n:100 7902281 ns 7229220 ns 95 svd/m:500/n:100 9720727 ns 9040679 ns 79 svd/m:800/n:100 9856378 ns 8998050 ns 79 svd/m:1000/n:100 9721017 ns 9086414 ns 79 svd/m:1/n:500 371171 ns 334217 ns 2117 svd/m:2/n:500 449165 ns 411499 ns 1700 svd/m:5/n:500 620354 ns 581866 ns 1185 svd/m:10/n:500 892375 ns 847239 ns 833 svd/m:100/n:500 9564810 ns 8867540 ns 79 svd/m:500/n:500 111924035 ns 104078023 ns 7 svd/m:800/n:500 147777319 ns 142730412 ns 5 svd/m:1000/n:500 154205084 ns 149740209 ns 5 svd/m:1/n:800 372122 ns 334212 ns 2119 svd/m:2/n:800 456672 ns 419260 ns 1680 svd/m:5/n:800 691208 ns 626003 ns 1190 svd/m:10/n:800 1017694 ns 941480 ns 730 svd/m:100/n:800 9892683 ns 9091043 ns 76 svd/m:500/n:800 144134235 ns 139129722 ns 5 svd/m:800/n:800 342790246 ns 333299774 ns 2 svd/m:1000/n:800 432820082 ns 427978978 ns 2 svd/m:1/n:1000 372785 ns 335745 ns 1805 svd/m:2/n:1000 451946 ns 413341 ns 1668 svd/m:5/n:1000 618475 ns 577213 ns 1169 svd/m:10/n:1000 907729 ns 863335 ns 808 svd/m:100/n:1000 9868543 ns 9116870 ns 76 svd/m:500/n:1000 156777811 ns 152042065 ns 5 svd/m:800/n:1000 429704070 ns 424677592 ns 2 svd/m:1000/n:1000 654864311 ns 642693162 ns 1 After: ------------------------------------------------------------ Benchmark Time CPU Iterations ------------------------------------------------------------ svd/m:1/n:1 265980 ns 245433 ns 2791 svd/m:2/n:1 340203 ns 302783 ns 2288 svd/m:5/n:1 337807 ns 301916 ns 2286 svd/m:10/n:1 338064 ns 302441 ns 2297 svd/m:100/n:1 335444 ns 298440 ns 2327 svd/m:500/n:1 338025 ns 302096 ns 2272 svd/m:800/n:1 328382 ns 291740 ns 2252 svd/m:1000/n:1 397494 ns 310905 ns 2239 svd/m:1/n:2 310464 ns 274507 ns 2535 svd/m:2/n:2 319999 ns 284247 ns 2515 svd/m:5/n:2 373435 ns 335919 ns 2069 svd/m:10/n:2 376327 ns 339327 ns 2056 svd/m:100/n:2 385061 ns 349258 ns 2003 svd/m:500/n:2 392352 ns 355735 ns 1932 svd/m:800/n:2 410736 ns 370677 ns 1881 svd/m:1000/n:2 494326 ns 405603 ns 1721 svd/m:1/n:5 316735 ns 277292 ns 2538 svd/m:2/n:5 383748 ns 342218 ns 2077 svd/m:5/n:5 494204 ns 454309 ns 1476 svd/m:10/n:5 547017 ns 508184 ns 1371 svd/m:100/n:5 514537 ns 476761 ns 1460 svd/m:500/n:5 544656 ns 504877 ns 1381 svd/m:800/n:5 642590 ns 599314 ns 1159 svd/m:1000/n:5 706166 ns 621209 ns 1106 svd/m:1/n:10 310825 ns 274374 ns 2511 svd/m:2/n:10 381316 ns 344202 ns 2094 svd/m:5/n:10 565469 ns 526759 ns 1266 svd/m:10/n:10 576111 ns 537286 ns 1299 svd/m:100/n:10 653250 ns 613392 ns 1137 svd/m:500/n:10 690532 ns 645828 ns 1080 svd/m:800/n:10 763924 ns 723677 ns 959 svd/m:1000/n:10 940342 ns 855517 ns 818 svd/m:1/n:100 306134 ns 271533 ns 2526 svd/m:2/n:100 374680 ns 339298 ns 2071 svd/m:5/n:100 576926 ns 539062 ns 1228 svd/m:10/n:100 656806 ns 615171 ns 1123 svd/m:100/n:100 3295164 ns 3138621 ns 223 svd/m:500/n:100 4269347 ns 4166000 ns 168 svd/m:800/n:100 4656541 ns 4522247 ns 154 svd/m:1000/n:100 6479223 ns 6354578 ns 112 svd/m:1/n:500 329966 ns 289083 ns 2440 svd/m:2/n:500 407535 ns 366794 ns 1947 svd/m:5/n:500 567367 ns 522809 ns 1336 svd/m:10/n:500 712307 ns 657608 ns 1065 svd/m:100/n:500 4262986 ns 4169907 ns 167 svd/m:500/n:500 28824720 ns 28650258 ns 25 svd/m:800/n:500 29330139 ns 28677269 ns 25 svd/m:1000/n:500 30848037 ns 30089216 ns 23 svd/m:1/n:800 328620 ns 289181 ns 2329 svd/m:2/n:800 419052 ns 379483 ns 1876 svd/m:5/n:800 587366 ns 546979 ns 1269 svd/m:10/n:800 830762 ns 787923 ns 893 svd/m:100/n:800 4763633 ns 4595738 ns 152 svd/m:500/n:800 30447861 ns 29949714 ns 24 svd/m:800/n:800 94188958 ns 93488372 ns 8 svd/m:1000/n:800 94701529 ns 93394677 ns 7 svd/m:1/n:1000 351102 ns 313099 ns 2218 svd/m:2/n:1000 446543 ns 407807 ns 1708 svd/m:5/n:1000 661152 ns 616174 ns 1129 svd/m:10/n:1000 915743 ns 873397 ns 802 svd/m:100/n:1000 6434730 ns 6282779 ns 113 svd/m:500/n:1000 30244321 ns 29684290 ns 24 svd/m:800/n:1000 92727423 ns 91477078 ns 8 svd/m:1000/n:1000 169500709 ns 168358420 ns 4 PiperOrigin-RevId: 582041508
2023-11-13 12:03:36 -08:00
* Changes
Reverts ef9075159a67a2b94526b65e4a2c2904a4a49046 PiperOrigin-RevId: 582789416
2023-11-15 13:35:14 -08:00
* In preparation for adding distributed CPU support, JAX now treats CPU
devices identically to GPU and TPU devices, that is:
* `jax.devices()` includes all devices present in a distributed job, even
those not local to the current process. `jax.local_devices()` still only
includes devices local to the current process, so if the change to
`jax.devices()` breaks you, you most likely want to use
`jax.local_devices()` instead.
* CPU devices now receive a globally unique ID number within a distributed
job; previously CPU devices would receive a process-local ID number.
* The `process_index` of each CPU device will now match any GPU or TPU
devices within the same process; previously the `process_index` of a CPU
device was always 0.
Use a Jacobi SVD solver for unbatched SVDs up to 1024x1024 on NVIDIA GPUs. The unbatched Jacobi solver is faster for small-moderate matrices, and the unbatched kernel doesn't have size restrictions. Timings on T4 GPU: Before: ------------------------------------------------------------ Benchmark Time CPU Iterations ------------------------------------------------------------ svd/m:1/n:1 263587 ns 242274 ns 2780 svd/m:2/n:1 335561 ns 298238 ns 2303 svd/m:5/n:1 337784 ns 299841 ns 2304 svd/m:10/n:1 339184 ns 300703 ns 2311 svd/m:100/n:1 359826 ns 320088 ns 2159 svd/m:500/n:1 376124 ns 338660 ns 2076 svd/m:800/n:1 375779 ns 335590 ns 2060 svd/m:1000/n:1 419171 ns 341487 ns 2072 svd/m:1/n:2 307564 ns 270663 ns 2544 svd/m:2/n:2 320928 ns 283601 ns 2487 svd/m:5/n:2 377373 ns 344228 ns 2035 svd/m:10/n:2 380557 ns 349412 ns 1953 svd/m:100/n:2 435465 ns 403496 ns 1722 svd/m:500/n:2 444610 ns 410913 ns 1680 svd/m:800/n:2 454493 ns 416495 ns 1665 svd/m:1000/n:2 492110 ns 420539 ns 1665 svd/m:1/n:5 307316 ns 275833 ns 2531 svd/m:2/n:5 374318 ns 341432 ns 2086 svd/m:5/n:5 512928 ns 470293 ns 1361 svd/m:10/n:5 589330 ns 537070 ns 1353 svd/m:100/n:5 620164 ns 580166 ns 1193 svd/m:500/n:5 636424 ns 593692 ns 1180 svd/m:800/n:5 635545 ns 595016 ns 1181 svd/m:1000/n:5 672443 ns 597387 ns 1115 svd/m:1/n:10 310013 ns 273998 ns 2520 svd/m:2/n:10 370451 ns 334489 ns 2105 svd/m:5/n:10 560037 ns 522223 ns 1274 svd/m:10/n:10 572868 ns 535388 ns 1304 svd/m:100/n:10 959802 ns 918258 ns 765 svd/m:500/n:10 955958 ns 909778 ns 758 svd/m:800/n:10 924104 ns 879512 ns 777 svd/m:1000/n:10 950140 ns 883493 ns 775 svd/m:1/n:100 351237 ns 315554 ns 2198 svd/m:2/n:100 426883 ns 390089 ns 1792 svd/m:5/n:100 601557 ns 564493 ns 1255 svd/m:10/n:100 920819 ns 880011 ns 787 svd/m:100/n:100 7902281 ns 7229220 ns 95 svd/m:500/n:100 9720727 ns 9040679 ns 79 svd/m:800/n:100 9856378 ns 8998050 ns 79 svd/m:1000/n:100 9721017 ns 9086414 ns 79 svd/m:1/n:500 371171 ns 334217 ns 2117 svd/m:2/n:500 449165 ns 411499 ns 1700 svd/m:5/n:500 620354 ns 581866 ns 1185 svd/m:10/n:500 892375 ns 847239 ns 833 svd/m:100/n:500 9564810 ns 8867540 ns 79 svd/m:500/n:500 111924035 ns 104078023 ns 7 svd/m:800/n:500 147777319 ns 142730412 ns 5 svd/m:1000/n:500 154205084 ns 149740209 ns 5 svd/m:1/n:800 372122 ns 334212 ns 2119 svd/m:2/n:800 456672 ns 419260 ns 1680 svd/m:5/n:800 691208 ns 626003 ns 1190 svd/m:10/n:800 1017694 ns 941480 ns 730 svd/m:100/n:800 9892683 ns 9091043 ns 76 svd/m:500/n:800 144134235 ns 139129722 ns 5 svd/m:800/n:800 342790246 ns 333299774 ns 2 svd/m:1000/n:800 432820082 ns 427978978 ns 2 svd/m:1/n:1000 372785 ns 335745 ns 1805 svd/m:2/n:1000 451946 ns 413341 ns 1668 svd/m:5/n:1000 618475 ns 577213 ns 1169 svd/m:10/n:1000 907729 ns 863335 ns 808 svd/m:100/n:1000 9868543 ns 9116870 ns 76 svd/m:500/n:1000 156777811 ns 152042065 ns 5 svd/m:800/n:1000 429704070 ns 424677592 ns 2 svd/m:1000/n:1000 654864311 ns 642693162 ns 1 After: ------------------------------------------------------------ Benchmark Time CPU Iterations ------------------------------------------------------------ svd/m:1/n:1 265980 ns 245433 ns 2791 svd/m:2/n:1 340203 ns 302783 ns 2288 svd/m:5/n:1 337807 ns 301916 ns 2286 svd/m:10/n:1 338064 ns 302441 ns 2297 svd/m:100/n:1 335444 ns 298440 ns 2327 svd/m:500/n:1 338025 ns 302096 ns 2272 svd/m:800/n:1 328382 ns 291740 ns 2252 svd/m:1000/n:1 397494 ns 310905 ns 2239 svd/m:1/n:2 310464 ns 274507 ns 2535 svd/m:2/n:2 319999 ns 284247 ns 2515 svd/m:5/n:2 373435 ns 335919 ns 2069 svd/m:10/n:2 376327 ns 339327 ns 2056 svd/m:100/n:2 385061 ns 349258 ns 2003 svd/m:500/n:2 392352 ns 355735 ns 1932 svd/m:800/n:2 410736 ns 370677 ns 1881 svd/m:1000/n:2 494326 ns 405603 ns 1721 svd/m:1/n:5 316735 ns 277292 ns 2538 svd/m:2/n:5 383748 ns 342218 ns 2077 svd/m:5/n:5 494204 ns 454309 ns 1476 svd/m:10/n:5 547017 ns 508184 ns 1371 svd/m:100/n:5 514537 ns 476761 ns 1460 svd/m:500/n:5 544656 ns 504877 ns 1381 svd/m:800/n:5 642590 ns 599314 ns 1159 svd/m:1000/n:5 706166 ns 621209 ns 1106 svd/m:1/n:10 310825 ns 274374 ns 2511 svd/m:2/n:10 381316 ns 344202 ns 2094 svd/m:5/n:10 565469 ns 526759 ns 1266 svd/m:10/n:10 576111 ns 537286 ns 1299 svd/m:100/n:10 653250 ns 613392 ns 1137 svd/m:500/n:10 690532 ns 645828 ns 1080 svd/m:800/n:10 763924 ns 723677 ns 959 svd/m:1000/n:10 940342 ns 855517 ns 818 svd/m:1/n:100 306134 ns 271533 ns 2526 svd/m:2/n:100 374680 ns 339298 ns 2071 svd/m:5/n:100 576926 ns 539062 ns 1228 svd/m:10/n:100 656806 ns 615171 ns 1123 svd/m:100/n:100 3295164 ns 3138621 ns 223 svd/m:500/n:100 4269347 ns 4166000 ns 168 svd/m:800/n:100 4656541 ns 4522247 ns 154 svd/m:1000/n:100 6479223 ns 6354578 ns 112 svd/m:1/n:500 329966 ns 289083 ns 2440 svd/m:2/n:500 407535 ns 366794 ns 1947 svd/m:5/n:500 567367 ns 522809 ns 1336 svd/m:10/n:500 712307 ns 657608 ns 1065 svd/m:100/n:500 4262986 ns 4169907 ns 167 svd/m:500/n:500 28824720 ns 28650258 ns 25 svd/m:800/n:500 29330139 ns 28677269 ns 25 svd/m:1000/n:500 30848037 ns 30089216 ns 23 svd/m:1/n:800 328620 ns 289181 ns 2329 svd/m:2/n:800 419052 ns 379483 ns 1876 svd/m:5/n:800 587366 ns 546979 ns 1269 svd/m:10/n:800 830762 ns 787923 ns 893 svd/m:100/n:800 4763633 ns 4595738 ns 152 svd/m:500/n:800 30447861 ns 29949714 ns 24 svd/m:800/n:800 94188958 ns 93488372 ns 8 svd/m:1000/n:800 94701529 ns 93394677 ns 7 svd/m:1/n:1000 351102 ns 313099 ns 2218 svd/m:2/n:1000 446543 ns 407807 ns 1708 svd/m:5/n:1000 661152 ns 616174 ns 1129 svd/m:10/n:1000 915743 ns 873397 ns 802 svd/m:100/n:1000 6434730 ns 6282779 ns 113 svd/m:500/n:1000 30244321 ns 29684290 ns 24 svd/m:800/n:1000 92727423 ns 91477078 ns 8 svd/m:1000/n:1000 169500709 ns 168358420 ns 4 PiperOrigin-RevId: 582041508
2023-11-13 12:03:36 -08:00
* On NVIDIA GPU, JAX now prefers a Jacobi SVD solver for matrices up to
1024x1024. The Jacobi solver appears faster than the non-Jacobi version.
Fix error/hang when non-finite values are passed to non-symmetric Eigendecomposition. Improve the documentation of lax.eig(). Fixes https://github.com/google/jax/issues/18226 PiperOrigin-RevId: 584170564
2023-11-20 17:27:42 -08:00
* Bug fixes
* Fixed error/hang when an array with non-finite values is passed to a
non-symmetric eigendecomposition (#18226). Arrays with non-finite values now
produce arrays full of NaNs as outputs.
Reverts ef9075159a67a2b94526b65e4a2c2904a4a49046 PiperOrigin-RevId: 582789416
2023-11-15 13:35:14 -08:00
Update versions and CHANGELOG after jax 0.4.20 release
2023-11-02 16:30:56 -07:00
## jax 0.4.20 (Nov 2, 2023)
## jaxlib 0.4.20 (Nov 2, 2023)
Finish 0.4.19 jax and jaxlib release PiperOrigin-RevId: 574983871
2023-10-19 13:27:16 -07:00
[IFRT] Fix incorrect type numbers for e4m3 and e5m2 types. These types didn't match between xla::PrimitiveType and ifrt::DType. Add a static_assert to enforce equality. PiperOrigin-RevId: 576288042
2023-10-24 14:37:28 -07:00
* Bug fixes
* Fixed some type confusion between E4M3 and E5M2 float8 types.
update change log's formatting. PiperOrigin-RevId: 576214943
2023-10-24 11:51:53 -07:00
## jax 0.4.19 (Oct 19, 2023)
Update versions and CHANGELOG after jax 0.4.17 release
2023-10-03 17:54:35 -07:00
Define jax.typing.DTypeLike
2023-10-10 08:46:36 -07:00
* New Features
* Added {obj}`jax.typing.DTypeLike`, which can be used to annotate objects that
are convertible to JAX dtypes.
Add jax.numpy.fill_diagonal.
2023-10-20 16:47:46 -04:00
* Added `jax.numpy.fill_diagonal`.
Define jax.typing.DTypeLike
2023-10-10 08:46:36 -07:00
Only perform compilation cache writes from process 0. This avoids problems with contending writes on filesystems such as GCS. PiperOrigin-RevId: 572032482
2023-10-09 13:54:32 -07:00
* Changes
* JAX now requires SciPy 1.9 or newer.
* Bug fixes
* Only process 0 in a multicontroller distributed JAX program will write
persistent compilation cache entries. This fixes write contention if the
Fix doc typos
2024-04-05 14:21:33 +05:30
cache is placed on a network file system such as GCS.
Only perform compilation cache writes from process 0. This avoids problems with contending writes on filesystems such as GCS. PiperOrigin-RevId: 572032482
2023-10-09 13:54:32 -07:00
* The version check for cusolver and cufft no longer considers the patch
versions when determining if the installed version of these libraries is at
least as new as the versions against which JAX was built.
update change log's formatting. PiperOrigin-RevId: 576214943
2023-10-24 11:51:53 -07:00
## jaxlib 0.4.19 (Oct 19, 2023)
Update versions post jax 0.4.18 release
2023-10-06 17:20:34 -07:00
Set RPATH, not RUNPATH in JAX CUDA builds. Fixes https://github.com/google/jax/issues/17497
2023-10-10 10:57:56 -04:00
* Changes
* jaxlib will now always prefer pip-installed NVIDIA CUDA libraries
(nvidia-... packages) over any other CUDA installation if they are
installed, including installations named in `LD_LIBRARY_PATH`. If this
causes problems and the intent is to use a system-installed CUDA, the fix is
to remove the pip installed CUDA library packages.
update change log's formatting. PiperOrigin-RevId: 576214943
2023-10-24 11:51:53 -07:00
## jax 0.4.18 (Oct 6, 2023)
Update versions post jax 0.4.18 release
2023-10-06 17:20:34 -07:00
update change log's formatting. PiperOrigin-RevId: 576214943
2023-10-24 11:51:53 -07:00
## jaxlib 0.4.18 (Oct 6, 2023)
Update versions and CHANGELOG after jax 0.4.17 release
2023-10-03 17:54:35 -07:00
[JAX] Obtain NCCL via a stub, rather than linking it statically or dynamically. This shrinks the CUDA jaxlib wheel size by around 80MB. PiperOrigin-RevId: 570554454
2023-10-03 18:33:09 -07:00
* Changes
* CUDA jaxlibs now depend on the user to install a compatible NCCL version.
If using the recommended `cuda12_pip` installation, NCCL should be installed
Improve CUDA install documentation. Mention NCCL as a dependency, since it will be required by the next jaxlib release. Mention LD_LIBRARY_PATH and PATH as how one overrides the CUDA installation for local installs. Fixes #17831
2023-10-06 09:11:49 -04:00
automatically. Currently, NCCL 2.16 or newer is required.
* We now provide Linux aarch64 wheels, both with and without NVIDIA GPU
support.
jnp.ndarray.item(): add args support
2024-01-03 13:03:47 -08:00
* {meth}`jax.Array.item` now supports optional index arguments.
[JAX] Obtain NCCL via a stub, rather than linking it statically or dynamically. This shrinks the CUDA jaxlib wheel size by around 80MB. PiperOrigin-RevId: 570554454
2023-10-03 18:33:09 -07:00
jax.lax: deprecate inadvertent exports & internal utilities
2023-10-06 11:26:03 -07:00
* Deprecations
* A number of internal utilities and inadvertent exports in {mod}`jax.lax` have
been deprecated, and will be removed in a future release.
* `jax.lax.dtypes`: use `jax.dtypes` instead.
* `jax.lax.itertools`: use `itertools` instead.
* `naryop`, `naryop_dtype_rule`, `standard_abstract_eval`, `standard_naryop`,
`standard_primitive`, `standard_unop`, `unop`, and `unop_dtype_rule` are
internal utilities, now deprecated without replacement.
Update setup.py and CHANGELOG for jax 0.4.18 release
2023-10-06 13:10:59 -07:00
* Bug fixes
* Fixed Cloud TPU regression where compilation would OOM due to smem.
update change log's formatting. PiperOrigin-RevId: 576214943
2023-10-24 11:51:53 -07:00
## jax 0.4.17 (Oct 3, 2023)
Update versions and changelog post 0.4.15 release
2023-08-30 15:31:01 -04:00
Add jax.numpy.bitwise_count()
2023-10-03 13:48:16 -07:00
* New features
Fix Typos in CHANGELOG.md This PR fixes the typos in Change log documentation
2024-03-12 13:57:07 +05:30
* Added new {func}`jax.numpy.bitwise_count` function, matching the API of the similar
Add jax.numpy.bitwise_count()
2023-10-03 13:48:16 -07:00
function recently added to NumPy.
Remove deprecated submodule jax.abstract_arrays
2023-09-19 15:40:18 -07:00
* Deprecations
* Removed the deprecated module `jax.abstract_arrays` and all its contents.
[random] deprecate named key creation functions
2023-09-21 13:57:49 -07:00
* Named key constructors in {mod}`jax.random` are deprecated. Pass the `impl` argument
to {func}`jax.random.PRNGKey` or {func}`jax.random.key` instead:
* `random.threefry2x32_key(seed)` becomes `random.PRNGKey(seed, impl='threefry2x32')`
* `random.rbg_key(seed)` becomes `random.PRNGKey(seed, impl='rbg')`
* `random.unsafe_rbg_key(seed)` becomes `random.PRNGKey(seed, impl='unsafe_rbg')`
Fix test failures under SciPy 1.11 for scipy.stats.mode.
2023-09-23 19:01:09 +00:00
* Changes:
Add notes about the new CUDA version restrictions to the changelog and installation instructions.
2023-09-27 15:56:47 -04:00
* CUDA: JAX now verifies that the CUDA libraries it finds are at least as new
as the CUDA libraries that JAX was built against. If older libraries are
found, JAX raises an exception since that is preferable to mysterious
failures and crashes.
Update changelog. Bump the minimum CUDA 12 pip package versions to the current releases.
2023-09-26 17:59:49 -04:00
* Removed the "No GPU/TPU" found warning. Instead warn if, on Linux, an
NVIDIA GPU or a Google TPU are found but not used and `--jax_platforms` was
not specified.
Add notes about the new CUDA version restrictions to the changelog and installation instructions.
2023-09-27 15:56:47 -04:00
* {func}`jax.scipy.stats.mode` now returns a 0 count if the mode is taken
across a size-0 axis, matching the behavior of `scipy.stats.mode` in SciPy
1.11.
Add CHANGELOG entry for the jnp annotation change
2023-10-02 11:31:28 -07:00
* Most `jax.numpy` functions and attributes now have fully-defined type stubs.
Previously many of these were treated as `Any` by static type checkers like
`mypy` and `pytype`.
Remove deprecated submodule jax.abstract_arrays
2023-09-19 15:40:18 -07:00
update change log's formatting. PiperOrigin-RevId: 576214943
2023-10-24 11:51:53 -07:00
## jaxlib 0.4.17 (Oct 3, 2023)
Add a release note about a fixed Windows crash.
2023-09-07 09:35:25 -04:00
Update changelog. Bump the minimum CUDA 12 pip package versions to the current releases.
2023-09-26 17:59:49 -04:00
* Changes:
* Python 3.12 wheels were added in this release.
Add notes about the new CUDA version restrictions to the changelog and installation instructions.
2023-09-27 15:56:47 -04:00
* The CUDA 12 wheels now require CUDA 12.2 or newer and cuDNN 8.9.4 or newer.
Update changelog. Bump the minimum CUDA 12 pip package versions to the current releases.
2023-09-26 17:59:49 -04:00
* Bug fixes:
* Fixed log spam from ABSL when the JAX CPU backend was initialized.
Finish jax and jaxlib 0.4.16 release PiperOrigin-RevId: 566477931
2023-09-18 19:08:43 -07:00
## jax 0.4.16 (Sept 18, 2023)
Update versions and changelog for jax 0.4.14 release
2023-07-27 16:22:53 -07:00
Documented the shortening of tracebacks
2023-08-04 11:11:05 -07:00
* Changes
Add jnp.ufunc and jnp.frompyfunc
2023-08-10 14:58:18 -07:00
* Added {class}`jax.numpy.ufunc`, as well as {func}`jax.numpy.frompyfunc`, which can convert
any scalar-valued function into a {func}`numpy.ufunc`-like object, with methods such as
{meth}`~jax.numpy.ufunc.outer`, {meth}`~jax.numpy.ufunc.reduce`,
{meth}`~jax.numpy.ufunc.accumulate`, {meth}`~jax.numpy.ufunc.at`, and
{meth}`~jax.numpy.ufunc.reduceat` ({jax-issue}`#17054`).
Deprecate jax.numpy.trapz. Expose the current implementation of jax.numpy.trapz as jax.scipy.integrate.trapezoid instead. Fixes https://github.com/google/jax/issues/17244
2023-08-24 14:01:40 -06:00
* Added {func}`jax.scipy.integrate.trapezoid`.
Documented the shortening of tracebacks
2023-08-04 11:11:05 -07:00
* When not running under IPython: when an exception is raised, JAX now filters out the
entirety of its internal frames from tracebacks. (Without the "unfiltered stack trace"
that previously appeared.) This should produce much friendlier-looking tracebacks. See
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
[here](https://github.com/jax-ml/jax/pull/16949) for an example.
Documented the shortening of tracebacks
2023-08-04 11:11:05 -07:00
This behavior can be changed by setting `JAX_TRACEBACK_FILTERING=remove_frames` (for two
separate unfiltered/filtered tracebacks, which was the old behavior) or
`JAX_TRACEBACK_FILTERING=off` (for one unfiltered traceback).
[jax2tf] Bump the default JAX serialization version to 7. This enables shape assertion checking, the support for which landed in XlaCallModule on July 12th, 2023. See the CHANGELOG for details. PiperOrigin-RevId: 556222908
2023-08-11 22:49:04 -07:00
* jax2tf default serialization version is now 7, which introduces new shape
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
[safety assertions](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/README.md#errors-in-presence-of-shape-polymorphism).
Any devices passed to jax.sharding.Mesh are required to be hashable. This is true for mock devices or user specific devices and jax.devices() too. Fix the tests so that the mock devices are hashable. PiperOrigin-RevId: 561103167
2023-08-29 12:17:37 -07:00
* Devices passed to `jax.sharding.Mesh` should be hashable. This specifically
applies to mock devices or user created devices. `jax.devices()` are
already hashable.
Documented the shortening of tracebacks
2023-08-04 11:11:05 -07:00
[jax2tf] Turn on JAX native serialization by default. See changes to the README.md for mechanisms to override the default. PiperOrigin-RevId: 554390866
2023-08-07 01:03:15 -07:00
* Breaking changes:
* jax2tf now uses native serialization by default. See
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
the [jax2tf documentation](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/README.md)
[jax2tf] Turn on JAX native serialization by default. See changes to the README.md for mechanisms to override the default. PiperOrigin-RevId: 554390866
2023-08-07 01:03:15 -07:00
for details and for mechanisms to override the default.
[JAX] Remove the non-coordination service distributed service implementation from JAX. The coordination service has been the default for a long time, and has significant additional functionality. Remove the older code path to simplify the code. PiperOrigin-RevId: 554608165
2023-08-07 15:16:50 -07:00
* The option `--jax_coordination_service` has been removed. It is now always
`True`.
Move jax.jaxpr_util to jax._src.jaxpr_util, and split it into a separate build target. Change jaxpr_util_test to be a py_test(), since there's no point testing it on every hardware configuration. PiperOrigin-RevId: 554861284
2023-08-08 10:08:19 -07:00
* `jax.jaxpr_util` has been removed from the public JAX namespace.
Remove option to use StreamExecutor Cloud TPU client in JAX It's been over three months since the new PJRT C API client was enabled by default (https://jax.readthedocs.io/en/latest/changelog.html#jax-0-4-8-march-29-2023). PiperOrigin-RevId: 554935166
2023-08-08 14:04:34 -07:00
* `JAX_USE_PJRT_C_API_ON_TPU` no longer has an effect (i.e. it always defaults to true).
[host_callback] Remove old backwards compatibility flag jax_host_callback_ad_transforms. This flag was added in https://github.com/google/jax/pull/8678 in December 2021 when we changed the behavior of host_callback to not have special handling for autodiff. Nobody is using that flag now. This is part of a longer project to replace uses of host_callback with jax.pure_callback and jax.experimental.io_callback. PiperOrigin-RevId: 557520668
2023-08-16 10:01:11 -07:00
* The backwards compatibility flag `--jax_host_callback_ad_transforms`
introduced in December 2021, has been removed.
[jax2tf] Turn on JAX native serialization by default. See changes to the README.md for mechanisms to override the default. PiperOrigin-RevId: 554390866
2023-08-07 01:03:15 -07:00
Mark jnp.{NINF,NZERO,PZERO} as deprecated. This follows the upstream NumPy deprecation of these names (https://github.com/numpy/numpy/pull/24357). PiperOrigin-RevId: 555548986
2023-08-10 10:24:40 -07:00
* Deprecations:
* Several `jax.numpy` APIs have been deprecated following
[NumPy NEP-52](https://numpy.org/neps/nep-0052-python-api-cleanup.html):
* `jax.numpy.NINF` has been deprecated. Use `-jax.numpy.inf` instead.
* `jax.numpy.PZERO` has been deprecated. Use `0.0` instead.
* `jax.numpy.NZERO` has been deprecated. Use `-0.0` instead.
Deprecate jax.numpy.row_stack
2023-08-22 13:12:49 -07:00
* `jax.numpy.issubsctype(x, t)` has been deprecated. Use `jax.numpy.issubdtype(x.dtype, t)`.
* `jax.numpy.row_stack` has been deprecated. Use `jax.numpy.vstack` instead.
Deprecate jax.numpy.in1d. Issue https://github.com/google/jax/issues/17244
2023-08-23 17:30:01 -06:00
* `jax.numpy.in1d` has been deprecated. Use `jax.numpy.isin` instead.
Deprecate jax.numpy.trapz. Expose the current implementation of jax.numpy.trapz as jax.scipy.integrate.trapezoid instead. Fixes https://github.com/google/jax/issues/17244
2023-08-24 14:01:40 -06:00
* `jax.numpy.trapz` has been deprecated. Use `jax.scipy.integrate.trapezoid` instead.
Deprecate jax.scipy.linalg.tril and jax.scipy.linalg.triu. The corresponding functions are deprecated in scipy. Use the equivalent jax.numpy functions instead.
2023-08-18 16:14:42 -04:00
* `jax.scipy.linalg.tril` and `jax.scipy.linalg.triu` have been deprecated,
following SciPy. Use `jax.numpy.tril` and `jax.numpy.triu` instead.
remove deprecated jax.lax.prod function PiperOrigin-RevId: 559787522
2023-08-24 10:02:27 -07:00
* `jax.lax.prod` has been removed after being deprecated in JAX v0.4.11.
Use the built-in `math.prod` instead.
Deprecate a number of exports from jax.interpreters.xla. Custom HLO lowering rules for primitives should be updated to use MLIR StableHLO lowering rules via jax.interpreter.mlir. PiperOrigin-RevId: 561215967
2023-08-29 20:47:22 -07:00
* A number of exports from `jax.interpreters.xla` related to defining
HLO lowering rules for custom JAX primitives have been deprecated. Custom
primitives should be defined using the StableHLO lowering utilities in
`jax.interpreters.mlir` instead.
Finish jax and jaxlib 0.4.16 release PiperOrigin-RevId: 566477931
2023-09-18 19:08:43 -07:00
* The following previously-deprecated functions have been removed after a
three-month deprecation period:
* `jax.abstract_arrays.ShapedArray`: use `jax.core.ShapedArray`.
* `jax.abstract_arrays.raise_to_shaped`: use `jax.core.raise_to_shaped`.
* `jax.numpy.alltrue`: use `jax.numpy.all`.
* `jax.numpy.sometrue`: use `jax.numpy.any`.
* `jax.numpy.product`: use `jax.numpy.prod`.
* `jax.numpy.cumproduct`: use `jax.numpy.cumprod`.
Mark jnp.{NINF,NZERO,PZERO} as deprecated. This follows the upstream NumPy deprecation of these names (https://github.com/numpy/numpy/pull/24357). PiperOrigin-RevId: 555548986
2023-08-10 10:24:40 -07:00
Finish jax and jaxlib 0.4.16 release PiperOrigin-RevId: 566477931
2023-09-18 19:08:43 -07:00
* Deprecations/removals:
* The internal submodule `jax.prng` is now deprecated. Its contents are available at
{mod}`jax.extend.random`.
* The internal submodule path `jax.linear_util` has been deprecated. Use
{mod}`jax.extend.linear_util` instead (Part of {ref}`jax-extend-jep`)
* `jax.random.PRNGKeyArray` and `jax.random.KeyArray` are deprecated. Use {class}`jax.Array`
dtypes.issubdtype: validate a when b is dtypes.extended
2023-09-21 11:20:54 -07:00
for type annotations, and `jax.dtypes.issubdtype(arr.dtype, jax.dtypes.prng_key)` for
runtime detection of typed prng keys.
Finish jax and jaxlib 0.4.16 release PiperOrigin-RevId: 566477931
2023-09-18 19:08:43 -07:00
* The method `PRNGKeyArray.unsafe_raw_array` is deprecated. Use
{func}`jax.random.key_data` instead.
* `jax.experimental.pjit.with_sharding_constraint` is deprecated. Use
`jax.lax.with_sharding_constraint` instead.
COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/17027 from jakevdp:dtypes-annotations a116a9c498a7b085f9b3fec93b37da12289f6e31 PiperOrigin-RevId: 554905739
2023-08-08 19:26:58 +00:00
* The internal utilities `jax.core.is_opaque_dtype` and `jax.core.has_opaque_dtype`
have been removed. Opaque dtypes have been renamed to Extended dtypes; use
`jnp.issubdtype(dtype, jax.dtypes.extended)` instead (available since jax v0.4.14).
Remove jax.interpreters.xla.register_collective_primitive. We aren't consuming this data any more. It existed only to compare against the set of multiprocess-allowed collectives, but we removed that list also. So this registry is completely pointless. PiperOrigin-RevId: 561150259
2023-08-29 15:09:33 -07:00
* The utility `jax.interpreters.xla.register_collective_primitive` has been
removed. This utility did nothing useful in recent JAX releases and calls
to it can be safely removed.
JEX: move jax.linear_util to jax.extend.linear_util
2023-08-30 15:14:47 -07:00
* The internal submodule path `jax.linear_util` has been deprecated. Use
{mod}`jax.extend.linear_util` instead (Part of {ref}`jax-extend-jep`)
COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/17027 from jakevdp:dtypes-annotations a116a9c498a7b085f9b3fec93b37da12289f6e31 PiperOrigin-RevId: 554905739
2023-08-08 19:26:58 +00:00
Finish jax and jaxlib 0.4.16 release PiperOrigin-RevId: 566477931
2023-09-18 19:08:43 -07:00
## jaxlib 0.4.16 (Sept 18, 2023)
Update versions and changelog for jax 0.4.14 release
2023-07-27 16:22:53 -07:00
Use the default CSR matmul algorithm. Previously we requested CUSPARSE_SPMM_CSR_ALG3 in an attempt to get deterministic results from cusparse SpMM CSR matmuls. In the past, Cusparse silently ignored this algorithm choice and used a different algorithm in cases where ALG3 was not supported, but cusparse 12.2.1 removed the silent fallback behavior. Since we're not actually getting deterministic behavior anyway in all cases, use the default algorithm always. PiperOrigin-RevId: 560867049
2023-08-28 17:48:19 -07:00
* Changes:
* Sparse CSR matrix multiplications via the experimental jax sparse APIs
no longer uses a deterministic algorithm on NVIDIA GPUs. This change was
made to improve compatibility with CUDA 12.2.1.
Finish jax and jaxlib 0.4.16 release PiperOrigin-RevId: 566477931
2023-09-18 19:08:43 -07:00
* Bug fixes:
* Fixed a crash on Windows due to a fatal LLVM error related to out-of-order
sections and IMAGE_REL_AMD64_ADDR32NB relocations
(https://github.com/openxla/xla/commit/cb732a921f0c4184995cbed82394931011d12bd4).
Update versions and changelog for jax 0.4.14 release
2023-07-27 16:22:53 -07:00
## jax 0.4.14 (July 27, 2023)
Update version numbers and changelog post 0.4.13 release
2023-06-22 17:54:57 -07:00
jax.jit now works correctly if both donate_argnums and donate_argnames are specified. Update the docstring and changelog too to mention `donate_argnames`. PiperOrigin-RevId: 548223395
2023-07-14 14:27:29 -07:00
* Changes
* `jax.jit` takes `donate_argnames` as an argument. It's semantics are similar
to `static_argnames`.
If neither donate_argnums nor donate_argnames is provided, no
arguments are donated. If donate_argnums is not provided but
donate_argnames is, or vice versa, JAX uses
`inspect.signature(fun)` to find any positional arguments that
correspond to donate_argnames (or vice versa). If both donate_argnums and donate_argnames are provided, inspect.signature is not used, and only actual
parameters listed in either donate_argnums or donate_argnames will
be donated.
Re-parameterize jax.random.gamma for better behavior at endpoints
2023-07-19 16:15:03 -07:00
* {func}`jax.random.gamma` has been re-factored to a more efficient algorithm
with more robust endpoint behavior ({jax-issue}`#16779`). This means that the
sequence of values returned for a given `key` will change between JAX v0.4.13
and v0.4.14 for `gamma` and related samplers (including {func}`jax.random.ball`,
Update versions and changelog for jax 0.4.14 release
2023-07-27 16:22:53 -07:00
{func}`jax.random.beta`, {func}`jax.random.chisquare`, {func}`jax.random.dirichlet`,
Re-parameterize jax.random.gamma for better behavior at endpoints
2023-07-19 16:15:03 -07:00
{func}`jax.random.generalized_normal`, {func}`jax.random.loggamma`, {func}`jax.random.t`).
jax.jit now works correctly if both donate_argnums and donate_argnames are specified. Update the docstring and changelog too to mention `donate_argnames`. PiperOrigin-RevId: 548223395
2023-07-14 14:27:29 -07:00
Delete in_axis_resources and out_axis_resources from pjit since it's been more than 3 months since their deprecation. The replace is to use in_shardings and out_shardings. You can still pass PartitionSpecs to {in|out}_shardings to pjit. PiperOrigin-RevId: 548673905
2023-07-17 06:35:13 -07:00
* Deletions
* `in_axis_resources` and `out_axis_resources` have been deleted from pjit since
it has been more than 3 months since their deprecation. Please use
`in_shardings` and `out_shardings` as the replacement.
This is a safe and trivial name replacement. It does not change any of the
current pjit semantics and doesn't break any code.
You can still pass in `PartitionSpecs` to in_shardings and out_shardings.
[Rollback] Remove py3.8 support from jax as per https://jax.readthedocs.io/en/latest/deprecation.html PiperOrigin-RevId: 542724110
2023-06-22 18:29:02 -07:00
* Deprecations
* Python 3.8 support has been dropped as per
https://jax.readthedocs.io/en/latest/deprecation.html
Drop support for numpy 1.21 This is in accordance with NEP 29 and https://jax.readthedocs.io/en/latest/deprecation.html
2023-06-23 01:22:49 -07:00
* JAX now requires NumPy 1.22 or newer as per
https://jax.readthedocs.io/en/latest/deprecation.html
Require index update optional arguments to be passed by keyword. Passing these keywords by position has been deprecated and has raised a warning since JAX v0.4.7 (Released 27 March 2023) PiperOrigin-RevId: 544620172
2023-06-30 04:29:43 -07:00
* Passing optional arguments to {func}`jax.numpy.ndarray.at` by position is
no longer supported, after being deprecated in JAX version 0.4.7.
For example, instead of `x.at[i].get(True)`, use `x.at[i].get(indices_are_sorted=True)`
Remove several deprecated jax.Array methods: - arr.broadcast - arr.broadcast_in_dim - arr.split These have been deprecated since JAX v0.4.5 PiperOrigin-RevId: 547228974
2023-07-11 10:26:33 -07:00
* The following `jax.Array` methods have been removed, after being deprecated
in JAX v0.4.5:
* `jax.Array.broadcast`: use {func}`jax.lax.broadcast` instead.
* `jax.Array.broadcast_in_dim`: use {func}`jax.lax.broadcast_in_dim` instead.
* `jax.Array.split`: use {func}`jax.numpy.split` instead.
Remove several deprecated APIs
2023-07-11 12:42:32 -07:00
* The following APIs have been removed after previous deprecation:
* `jax.ad`: use {mod}`jax.interpreters.ad`.
* `jax.curry`: use ``curry = lambda f: partial(partial, f)``.
* `jax.partial_eval`: use {mod}`jax.interpreters.partial_eval`.
* `jax.pxla`: use {mod}`jax.interpreters.pxla`.
* `jax.xla`: use {mod}`jax.interpreters.xla`.
* `jax.ShapedArray`: use {class}`jax.core.ShapedArray`.
* `jax.interpreters.pxla.device_put`: use {func}`jax.device_put`.
* `jax.interpreters.pxla.make_sharded_device_array`: use {func}`jax.make_array_from_single_device_arrays`.
* `jax.interpreters.pxla.ShardedDeviceArray`: use {class}`jax.Array`.
* `jax.numpy.DeviceArray`: use {class}`jax.Array`.
Add deprecation of `jax.stages.Compiled.compiler_ir` to the change log PiperOrigin-RevId: 549415191
2023-07-19 13:48:22 -07:00
* `jax.stages.Compiled.compiler_ir`: use {func}`jax.stages.Compiled.as_text`.
[Rollback] Remove py3.8 support from jax as per https://jax.readthedocs.io/en/latest/deprecation.html PiperOrigin-RevId: 542724110
2023-06-22 18:29:02 -07:00
add corner-case `cond` resolution fix to changelog
2023-06-28 10:09:10 -07:00
* Breaking changes
Require ml_dtypes>=0.2
2023-07-07 12:07:44 -07:00
* JAX now requires ml_dtypes version 0.2.0 or newer.
add corner-case `cond` resolution fix to changelog
2023-06-28 10:09:10 -07:00
* To fix a corner case, calls to {func}`jax.lax.cond` with five
arguments will always resolve to the "common operands" `cond`
behavior (as documented) if the second and third arguments are
callable, even if other operands are callable as well. See
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
[#16413](https://github.com/jax-ml/jax/issues/16413).
Remove the jax_array config option, which does nothing. PiperOrigin-RevId: 548981491
2023-07-18 06:15:24 -07:00
* The deprecated config options `jax_array` and `jax_jit_pjit_api_merge`,
which did nothing, have been removed. These options have been true by
default for many releases.
add corner-case `cond` resolution fix to changelog
2023-06-28 10:09:10 -07:00
Copybara import of the project: -- 06bf5fe7b2ac97156df541bab989dc5beb1aff0c by George Necula <gcnecula@gmail.com>: [jax2tf] Added a flag and environment variable to control the serialization version. This allows us to control the serialization version to be compatible with the deployed version of tf.XlaCallModule. In particular, we can run most tests with the maximum available version, while keeping the default lower. COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/16746 from gnecula:tf_version 06bf5fe7b2ac97156df541bab989dc5beb1aff0c PiperOrigin-RevId: 548504243
2023-07-16 09:26:27 -07:00
* New features
* JAX now supports a configuration flag --jax_serialization_version
and a JAX_SERIALIZATION_VERSION environment variable to control the
serialization version ({jax-issue}`#16746`).
[shape_poly] Reimplement the shape constraint checking using shape assertions. Most of the functionality is for the JAX native serialization case. This relies on newly added functionality to xla_extension.refine_polymorphic_shapes that handles custom calls @static_assertion. As a beneficial side-effect now we get shape constraint checking for jax2tf graph serialization when the resulting function is executed in graph mode.
2023-07-03 17:31:31 +03:00
* jax2tf in presence of shape polymorphism now generates code that checks
certain shape constraints, if the serialization version is at least 7.
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
See https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/README.md#errors-in-presence-of-shape-polymorphism.
Copybara import of the project: -- 06bf5fe7b2ac97156df541bab989dc5beb1aff0c by George Necula <gcnecula@gmail.com>: [jax2tf] Added a flag and environment variable to control the serialization version. This allows us to control the serialization version to be compatible with the deployed version of tf.XlaCallModule. In particular, we can run most tests with the maximum available version, while keeping the default lower. COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/16746 from gnecula:tf_version 06bf5fe7b2ac97156df541bab989dc5beb1aff0c PiperOrigin-RevId: 548504243
2023-07-16 09:26:27 -07:00
Update versions and changelog for jax 0.4.14 release
2023-07-27 16:22:53 -07:00
## jaxlib 0.4.14 (July 27, 2023)
Update version numbers and changelog post 0.4.13 release
2023-06-22 17:54:57 -07:00
[Rollback] Remove py3.8 support from jax as per https://jax.readthedocs.io/en/latest/deprecation.html PiperOrigin-RevId: 542724110
2023-06-22 18:29:02 -07:00
* Deprecations
* Python 3.8 support has been dropped as per
https://jax.readthedocs.io/en/latest/deprecation.html
Update version numbers and changelog post 0.4.13 release
2023-06-22 17:54:57 -07:00
## jax 0.4.13 (June 22, 2023)
Allow None to be passed to in_shardings and out_shardings. The default is still UNSPECIFIED to handle edge cases around the old semantics where None is treated as fully replicated. The semantics are as follow: * if the mesh context manager is not provided, None will be treated as UNSPECIFIED for both in_shardings and out_shardings * If the mesh context manager is provided, None will be treated as fully replicated as per the old semantics. This will make sure that we don't break existing code depending on None meaning replicated but also start making the transition to None meaning UNSPECIFIED for jit and pjit. PiperOrigin-RevId: 540705660
2023-06-15 15:21:36 -07:00
* 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.
Jax 0.4.13 release.
2023-06-22 16:30:04 -04:00
* 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`.
Raise an error if a user passes None to `host_local_array_to_global_array` or `global_array_to_host_local_array` PiperOrigin-RevId: 543596009
2023-06-26 18:15:04 -07:00
* `None` is not a valid input to
`jax.experimental.multihost_utils.host_local_array_to_global_array` or
`jax.experimental.multihost_utils.global_array_to_host_local_array`.
Please use `jax.sharding.PartitionSpec()` if you wanted to replicate your
input.
Allow None to be passed to in_shardings and out_shardings. The default is still UNSPECIFIED to handle edge cases around the old semantics where None is treated as fully replicated. The semantics are as follow: * if the mesh context manager is not provided, None will be treated as UNSPECIFIED for both in_shardings and out_shardings * If the mesh context manager is provided, None will be treated as fully replicated as per the old semantics. This will make sure that we don't break existing code depending on None meaning replicated but also start making the transition to None meaning UNSPECIFIED for jit and pjit. PiperOrigin-RevId: 540705660
2023-06-15 15:21:36 -07:00
Fix cuda12 pip install. The wheel is now called cudnn89. Fixes #16362
2023-06-12 15:31:38 -04:00
* Bug fixes
* Fixed incorrect wheel name in CUDA 12 releases (#16362); the correct wheel
is named `cudnn89` instead of `cudnn88`.
Update versions and changelog post 0.4.12 release
2023-06-08 16:53:22 -07:00
[jax2tf] Add native_lowering_disabled_checks parameter to jax2tf.convert. Previously, we had a boolean `native_serialization_strict_checks` parameter that was disabling all safety checks. This mechanism had several disadvantages: * the mechanism did not differentiate between different safety checks. E.g., in order to disable checking of the custom call targets, one had to disable checking for all custom call targets, and also the checking that the serialization and execution platforms are the same. * the mechanism operated only at serialization time. Now, the XlaCallModule supports a `disabled_checks` attribute to control which safety checks should be disabled. Here we replace the `native_serialization_strict_checks` with `native_serialization_disabled_checks`, whose values are sequences of disabled check descriptors.
2023-06-10 09:27:42 +03:00
* 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`).
Update version numbers and changelog post 0.4.13 release
2023-06-22 17:54:57 -07:00
## jaxlib 0.4.13 (June 22, 2023)
Update versions and changelog post 0.4.12 release
2023-06-08 16:53:22 -07:00
Jax 0.4.13 release.
2023-06-22 16:30:04 -04:00
* Changes
* Added Windows CPU-only wheels to the `jaxlib` Pypi release.
[XLA:Python] Fix __cuda_array_interface__. Adds a test for __cuda_array_interface__ that does not depend on cupy. Fixes https://github.com/google/jax/issues/16440 PiperOrigin-RevId: 541965361
2023-06-20 10:08:38 -07:00
* Bug fixes
* `__cuda_array_interface__` was broken in previous jaxlib versions and is now
fixed ({jax-issue}`16440`).
Jax 0.4.13 release.
2023-06-22 16:30:04 -04:00
* Concurrent CUDA kernel tracing is now enabled by default on NVIDIA GPUs.
[XLA:Python] Fix __cuda_array_interface__. Adds a test for __cuda_array_interface__ that does not depend on cupy. Fixes https://github.com/google/jax/issues/16440 PiperOrigin-RevId: 541965361
2023-06-20 10:08:38 -07:00
Update versions and changelog post 0.4.12 release
2023-06-08 16:53:22 -07:00
## jax 0.4.12 (June 8, 2023)
Finish jax and jaxlib 0.4.11 release PiperOrigin-RevId: 536931532
2023-05-31 23:48:50 -07:00
Added scipy.spatial.transform Rotation and Slerp classes
2023-06-02 22:38:15 -05:00
* Changes
* Added {class}`scipy.spatial.transform.Rotation` and {class}`scipy.spatial.transform.Slerp`
Deprecate jax.numpy functions alltrue, sometrue, product, cumproduct
2023-06-02 04:10:46 -07:00
* Deprecations
Mark jax.abstract_arrays as deprecated
2023-06-06 07:32:35 -07:00
* `jax.abstract_arrays` and its contents are now deprecated. See related
functionality in :mod:`jax.core`.
* `jax.numpy.alltrue`: use `jax.numpy.all`. This follows the deprecation
of `numpy.alltrue` in NumPy version 1.25.0.
* `jax.numpy.sometrue`: use `jax.numpy.any`. This follows the deprecation
of `numpy.sometrue` in NumPy version 1.25.0.
* `jax.numpy.product`: use `jax.numpy.prod`. This follows the deprecation
of `numpy.product` in NumPy version 1.25.0.
* `jax.numpy.cumproduct`: use `jax.numpy.cumprod`. This follows the deprecation
of `numpy.cumproduct` in NumPy version 1.25.0.
Delete OpShardingSharding export since it has been 3 months since it was deprecated. Also remove deprecation warnings for MeshPspecSharding. PiperOrigin-RevId: 538880293
2023-06-08 13:45:43 -07:00
* `jax.sharding.OpShardingSharding` has been removed since it has been 3
months since it was deprecated.
Deprecate jax.numpy functions alltrue, sometrue, product, cumproduct
2023-06-02 04:10:46 -07:00
Update versions and changelog post 0.4.12 release
2023-06-08 16:53:22 -07:00
## jaxlib 0.4.12 (June 8, 2023)
Finish jax and jaxlib 0.4.11 release PiperOrigin-RevId: 536931532
2023-05-31 23:48:50 -07:00
Include SASS/PTX for Hopper GPUs.
2023-06-05 09:42:12 -04:00
* Changes
[XLA:Python] Fix incorrect code source information under Python 3.11. Do not multiply the result of PyFrame_GetLasti() by sizeof(_Py_CODEUNIT), because the CPython implementation already does this inside PyFrame_GetLasti(). * In CPython versions 3.9 or earlier, the f_lasti value in a PyFrameObject was in bytes. * In CPython 3.10, f_lasti was changed to be in code units, which required multiplying it by sizeof(_Py_CODEUNIT) before passing it to functions like PyCode_Addr2Line(). https://docs.python.org/3/whatsnew/3.10.html#changes-in-the-c-api * In CPython 3.11, direct access to the representation of the PyFrameObject was removed from the headers, requiring the use of PyFrame_GetLasti() (https://docs.python.org/3/whatsnew/3.11.html#pyframeobject-3-11-hiding). This function multiplies by sizeof(_Py_CODEUNIT) internally (https://github.com/python/cpython/blob/deaf509e8fc6e0363bd6f26d52ad42f976ec42f2/Objects/frameobject.c#L1353) so there is no need for the caller to do this multiplication. It is difficult to write a good test for this, since the only symptom is slightly inaccurate code line information. This issue was found under a debug mode build of CPython (https://docs.python.org/3/using/configure.html#python-debug-build), where PyCode_Addr2Line() has additional checks for out of range lasti values. PiperOrigin-RevId: 538847288
2023-06-08 11:41:47 -07:00
* Includes PTX/SASS for Hopper (SM version 9.0+) GPUs. Previous
Include SASS/PTX for Hopper GPUs.
2023-06-05 09:42:12 -04:00
versions of jaxlib should work on Hopper but would have a long
JIT-compilation delay the first time a JAX operation was executed.
[XLA:Python] Fix incorrect code source information under Python 3.11. Do not multiply the result of PyFrame_GetLasti() by sizeof(_Py_CODEUNIT), because the CPython implementation already does this inside PyFrame_GetLasti(). * In CPython versions 3.9 or earlier, the f_lasti value in a PyFrameObject was in bytes. * In CPython 3.10, f_lasti was changed to be in code units, which required multiplying it by sizeof(_Py_CODEUNIT) before passing it to functions like PyCode_Addr2Line(). https://docs.python.org/3/whatsnew/3.10.html#changes-in-the-c-api * In CPython 3.11, direct access to the representation of the PyFrameObject was removed from the headers, requiring the use of PyFrame_GetLasti() (https://docs.python.org/3/whatsnew/3.11.html#pyframeobject-3-11-hiding). This function multiplies by sizeof(_Py_CODEUNIT) internally (https://github.com/python/cpython/blob/deaf509e8fc6e0363bd6f26d52ad42f976ec42f2/Objects/frameobject.c#L1353) so there is no need for the caller to do this multiplication. It is difficult to write a good test for this, since the only symptom is slightly inaccurate code line information. This issue was found under a debug mode build of CPython (https://docs.python.org/3/using/configure.html#python-debug-build), where PyCode_Addr2Line() has additional checks for out of range lasti values. PiperOrigin-RevId: 538847288
2023-06-08 11:41:47 -07:00
* Bug fixes
* Fixes incorrect source line information in JAX-generated Python tracebacks
under Python 3.11.
[XLA:Python] Fix crash when accessing locals of JAX-generated Python tracebacks. Use a dummy PyCodeObject in our dummy PyFrameObjects. Using a real PyCodeObject with a fake PyFrameObject confuses CPython 3.11+ when it attempts to compute the locals of a frame, since the frame lacks certain details such as closure information. This unfortunately means we will not get source column information under Python 3.11 any more, but that is probably better than crashing. Fixes https://github.com/google/jax/issues/16027 PiperOrigin-RevId: 538873850
2023-06-08 13:21:43 -07:00
* Fixes crash when printing local variables of frames in JAX-generated Python
tracebacks (#16027).
[XLA:Python] Fix incorrect code source information under Python 3.11. Do not multiply the result of PyFrame_GetLasti() by sizeof(_Py_CODEUNIT), because the CPython implementation already does this inside PyFrame_GetLasti(). * In CPython versions 3.9 or earlier, the f_lasti value in a PyFrameObject was in bytes. * In CPython 3.10, f_lasti was changed to be in code units, which required multiplying it by sizeof(_Py_CODEUNIT) before passing it to functions like PyCode_Addr2Line(). https://docs.python.org/3/whatsnew/3.10.html#changes-in-the-c-api * In CPython 3.11, direct access to the representation of the PyFrameObject was removed from the headers, requiring the use of PyFrame_GetLasti() (https://docs.python.org/3/whatsnew/3.11.html#pyframeobject-3-11-hiding). This function multiplies by sizeof(_Py_CODEUNIT) internally (https://github.com/python/cpython/blob/deaf509e8fc6e0363bd6f26d52ad42f976ec42f2/Objects/frameobject.c#L1353) so there is no need for the caller to do this multiplication. It is difficult to write a good test for this, since the only symptom is slightly inaccurate code line information. This issue was found under a debug mode build of CPython (https://docs.python.org/3/using/configure.html#python-debug-build), where PyCode_Addr2Line() has additional checks for out of range lasti values. PiperOrigin-RevId: 538847288
2023-06-08 11:41:47 -07:00
Finish jax and jaxlib 0.4.11 release PiperOrigin-RevId: 536931532
2023-05-31 23:48:50 -07:00
## jax 0.4.11 (May 31, 2023)
Update versions and changelog for 0.4.9 release
2023-05-09 17:06:59 -07:00
Copybara import of the project: -- 8cf6a6acd151007935b0c3093df05ef036bb0244 by Jake VanderPlas <jakevdp@google.com>: Remove several deprecated APIs COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/16110 from jakevdp:deprecations 8cf6a6acd151007935b0c3093df05ef036bb0244 PiperOrigin-RevId: 534897394
2023-05-24 10:35:01 -07:00
* Deprecations
* The following APIs have been removed after a 3 month deprecation period, in
accordance with the {ref}`api-compatibility` policy:
Remove axis_resources from with_sharding_constraint since it has been 3 months since the deprecation as per the API deprecation policy. PiperOrigin-RevId: 535687618
2023-05-26 12:34:32 -07:00
* `jax.experimental.PartitionSpec`: use `jax.sharding.PartitionSpec`.
* `jax.experimental.maps.Mesh`: use `jax.sharding.Mesh`
* `jax.experimental.pjit.NamedSharding`: use `jax.sharding.NamedSharding`.
* `jax.experimental.pjit.PartitionSpec`: use `jax.sharding.PartitionSpec`.
* `jax.experimental.pjit.FROM_GDA`. Instead pass sharded `jax.Array` objects
Copybara import of the project: -- 8cf6a6acd151007935b0c3093df05ef036bb0244 by Jake VanderPlas <jakevdp@google.com>: Remove several deprecated APIs COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/16110 from jakevdp:deprecations 8cf6a6acd151007935b0c3093df05ef036bb0244 PiperOrigin-RevId: 534897394
2023-05-24 10:35:01 -07:00
as input and remove the optional `in_shardings` argument to `pjit`.
Remove axis_resources from with_sharding_constraint since it has been 3 months since the deprecation as per the API deprecation policy. PiperOrigin-RevId: 535687618
2023-05-26 12:34:32 -07:00
* `jax.interpreters.pxla.PartitionSpec`: use `jax.sharding.PartitionSpec`.
* `jax.interpreters.pxla.Mesh`: use `jax.sharding.Mesh`
Deprecate jax.interpreters.xla.Buffer, device_put, xla_call_p
2023-05-28 07:15:34 -07:00
* `jax.interpreters.xla.Buffer`: use `jax.Array`.
Remove axis_resources from with_sharding_constraint since it has been 3 months since the deprecation as per the API deprecation policy. PiperOrigin-RevId: 535687618
2023-05-26 12:34:32 -07:00
* `jax.interpreters.xla.Device`: use `jax.Device`.
Deprecate jax.interpreters.xla.Buffer, device_put, xla_call_p
2023-05-28 07:15:34 -07:00
* `jax.interpreters.xla.DeviceArray`: use `jax.Array`.
* `jax.interpreters.xla.device_put`: use `jax.device_put`.
* `jax.interpreters.xla.xla_call_p`: use `jax.experimental.pjit.pjit_p`.
Remove axis_resources from with_sharding_constraint since it has been 3 months since the deprecation as per the API deprecation policy. PiperOrigin-RevId: 535687618
2023-05-26 12:34:32 -07:00
* `axis_resources` argument of `with_sharding_constraint` is removed. Please
use `shardings` instead.
Copybara import of the project: -- 8cf6a6acd151007935b0c3093df05ef036bb0244 by Jake VanderPlas <jakevdp@google.com>: Remove several deprecated APIs COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/16110 from jakevdp:deprecations 8cf6a6acd151007935b0c3093df05ef036bb0244 PiperOrigin-RevId: 534897394
2023-05-24 10:35:01 -07:00
Finish jax and jaxlib 0.4.11 release PiperOrigin-RevId: 536931532
2023-05-31 23:48:50 -07:00
## jaxlib 0.4.11 (May 31, 2023)
Update versions and changelog post 0.4.10 release
2023-05-11 18:16:02 -07:00
Reland: [XLA:Python] Add buffer protocol support to jax.Array We supported the buffer protocol on the older DeviceArray class; port that support to jax.Array. The previous attempt was reverted because it led to a C++ CHECK failure if the buffer was deleted while an external Python reference was held. Change the CPU PJRT client to keep the underlying buffer alive as long as there are external references, which is what the contract of Delete() says it will do. Fixes https://github.com/google/jax/issues/14713 PiperOrigin-RevId: 535248553
2023-05-25 07:19:56 -07:00
* Changes
[PJRT:C] Implement PJRT_Device_MemoryStats Inspired by https://github.com/google/jax/issues/1491#issuecomment-1567696310 PiperOrigin-RevId: 538596599
2023-06-07 14:41:34 -07:00
* Added `memory_stats()` method to `Device`s. If supported, this returns a
dict of string stat names with int values, e.g. `"bytes_in_use"`, or None if
the platform doesn't support memory statistics. The exact stats returned may
vary across platforms. Currently only implemented on Cloud TPU.
Reland: [XLA:Python] Add buffer protocol support to jax.Array We supported the buffer protocol on the older DeviceArray class; port that support to jax.Array. The previous attempt was reverted because it led to a C++ CHECK failure if the buffer was deleted while an external Python reference was held. Change the CPU PJRT client to keep the underlying buffer alive as long as there are external references, which is what the contract of Delete() says it will do. Fixes https://github.com/google/jax/issues/14713 PiperOrigin-RevId: 535248553
2023-05-25 07:19:56 -07:00
* Readded support for the Python buffer protocol (`memoryview`) on CPU
devices.
Update versions and changelog post 0.4.10 release
2023-05-11 18:16:02 -07:00
## jax 0.4.10 (May 11, 2023)
## jaxlib 0.4.10 (May 11, 2023)
Update versions and changelog for 0.4.9 release
2023-05-09 17:06:59 -07:00
Update setup.py, WORKSPACE, and CHANGELOG for jax/jaxlib 0.4.10 release
2023-05-11 14:46:06 -07:00
* Changes
* Fixed `'apple-m1' is not a recognized processor for this target (ignoring
processor)` issue that prevented previous release from running on Mac M1.
Update versions and changelog for 0.4.9 release
2023-05-09 17:06:59 -07:00
## jax 0.4.9 (May 9, 2023)
Update version and changelog after jax 0.4.8 release
2023-03-29 14:25:26 -07:00
Remove `experimental_cpp_jit` since that flag is unused and also remove `experimental_cpp_pjit`. For dynamic shapes experimentation and normal debugging, `python_pjit` still exists so that problem doesn't exist which makes us free to remove these 2 flags. I am leaving pmap's flag alone for now. PiperOrigin-RevId: 522602754
2023-04-07 08:28:46 -07:00
* Changes
Update versions and changelog for 0.4.9 release
2023-05-09 17:06:59 -07:00
* The flags experimental_cpp_jit, experimental_cpp_pjit and
Remove experimental_cpp_pmap flag since it is always on PiperOrigin-RevId: 522631405
2023-04-07 10:41:42 -07:00
experimental_cpp_pmap have been removed.
Remove `experimental_cpp_jit` since that flag is unused and also remove `experimental_cpp_pjit`. For dynamic shapes experimentation and normal debugging, `python_pjit` still exists so that problem doesn't exist which makes us free to remove these 2 flags. I am leaving pmap's flag alone for now. PiperOrigin-RevId: 522602754
2023-04-07 08:28:46 -07:00
They are now always on.
Add improved TPU SVD accuracy to the changelog. PiperOrigin-RevId: 530752990
2023-05-09 17:08:00 -07:00
* Accuracy of singular value decomposition (SVD) on TPU has been improved
(requires jaxlib 0.4.9).
Remove `experimental_cpp_jit` since that flag is unused and also remove `experimental_cpp_pjit`. For dynamic shapes experimentation and normal debugging, `python_pjit` still exists so that problem doesn't exist which makes us free to remove these 2 flags. I am leaving pmap's flag alone for now. PiperOrigin-RevId: 522602754
2023-04-07 08:28:46 -07:00
Raise deprecation warnings for `{in|out}_axis_resources` for pjit and `axis_resources` for with_sharding_constraint PiperOrigin-RevId: 520748845
2023-03-30 14:50:07 -07:00
* Deprecations
Rename gda_serialization to array_serialization but keep gda_serialization around until it is included in a jax release so that OSS projects can be moved to array_serialization PiperOrigin-RevId: 521055760
2023-03-31 18:07:08 -07:00
* `jax.experimental.gda_serialization` is deprecated and has been renamed to
`jax.experimental.array_serialization`.
Please change your imports to use `jax.experimental.array_serialization`.
Raise deprecation warnings for `{in|out}_axis_resources` for pjit and `axis_resources` for with_sharding_constraint PiperOrigin-RevId: 520748845
2023-03-30 14:50:07 -07:00
* The `in_axis_resources` and `out_axis_resources` arguments of pjit have been
deprecated. Please use `in_shardings` and `out_shardings` respectively.
Remove deprecated function jnp.msort
2023-03-30 13:18:28 -07:00
* The function `jax.numpy.msort` has been removed. It has been deprecated since
JAX v0.4.1. Use `jnp.sort(a, axis=0)` instead.
Remove in_parts, out_parts from jax.xla_computation since they were only used for sharded_jit and sharded_jit is long gone Also remove instantiate_const_outputs since that is unused PiperOrigin-RevId: 524113088
2023-04-13 15:04:48 -07:00
* `in_parts` and `out_parts` arguments have been removed from `jax.xla_computation`
since they were only used with sharded_jit and sharded_jit is long gone.
Deprecate and raise an exception for `instantiate_const_outputs` argument of jax.xla_computation since it has been unused for a very long time. PiperOrigin-RevId: 524295738
2023-04-14 08:19:31 -07:00
* `instantiate_const_outputs` argument has been removed from `jax.xla_computation`
since it has been unused for a very long time.
Raise deprecation warnings for `{in|out}_axis_resources` for pjit and `axis_resources` for with_sharding_constraint PiperOrigin-RevId: 520748845
2023-03-30 14:50:07 -07:00
Update versions and changelog for 0.4.9 release
2023-05-09 17:06:59 -07:00
## jaxlib 0.4.9 (May 9, 2023)
Update version and changelog after jax 0.4.8 release
2023-03-29 14:25:26 -07:00
## jax 0.4.8 (March 29, 2023)
Finish jax and jaxlib 0.4.7 release PiperOrigin-RevId: 519839723
2023-03-27 15:06:06 -07:00
Turn on PJRT C API by default. I forgot that the default setting is actually in jaxlib: https://github.com/openxla/xla/blob/fbe9a80fdb8c429e8a175962459da348cd560a50/xla/python/xla_client.py#L135 To be able to make this change as a jax-only release, I manually set the env var on Cloud TPU if it isn't already set.
2023-03-28 13:42:51 -07:00
* Breaking changes
* A major component of the Cloud TPU runtime has been upgraded. This enables
the following new features on Cloud TPU:
* {func}`jax.debug.print`, {func}`jax.debug.callback`, and
{func}`jax.debug.breakpoint()` now work on Cloud TPU
* Automatic TPU memory defragmentation
{func}`jax.experimental.host_callback` is no longer supported on Cloud TPU
with the new runtime component. Please file an issue on the [JAX issue
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
tracker](https://github.com/jax-ml/jax/issues) if the new `jax.debug` APIs
Turn on PJRT C API by default. I forgot that the default setting is actually in jaxlib: https://github.com/openxla/xla/blob/fbe9a80fdb8c429e8a175962459da348cd560a50/xla/python/xla_client.py#L135 To be able to make this change as a jax-only release, I manually set the env var on Cloud TPU if it isn't already set.
2023-03-28 13:42:51 -07:00
are insufficient for your use case.
The old runtime component will be available for at least the next three
months by setting the environment variable
`JAX_USE_PJRT_C_API_ON_TPU=false`. If you find you need to disable the new
runtime for any reason, please let us know on the [JAX issue
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
tracker](https://github.com/jax-ml/jax/issues).
Turn on PJRT C API by default. I forgot that the default setting is actually in jaxlib: https://github.com/openxla/xla/blob/fbe9a80fdb8c429e8a175962459da348cd560a50/xla/python/xla_client.py#L135 To be able to make this change as a jax-only release, I manually set the env var on Cloud TPU if it isn't already set.
2023-03-28 13:42:51 -07:00
Bump minimum jaxlib version from 0.4.6 to 0.4.7. Also removes a bunch of dead version guards (0.4.7 has xla_extension_version 144 and mlir_api_version 47)
2023-03-28 12:43:32 -07:00
* Changes
* The minimum jaxlib version has been bumped from 0.4.6 to 0.4.7.
Remove Cuda 11.4 support. JAX from 0.4.8 release will support cuda 11.8 and cuda 12 PiperOrigin-RevId: 520056811
2023-03-28 09:53:57 -07:00
* Deprecations
* CUDA 11.4 support has been dropped. JAX GPU wheels only support
CUDA 11.8 and CUDA 12. Older CUDA versions may work if jaxlib is built
from source.
Remove global_arg_shapes from pmap since it was only used for sharded_jit and sharded_jit was removed from JAX a long time ago PiperOrigin-RevId: 520356179
2023-03-29 09:22:34 -07:00
* `global_arg_shapes` argument of pmap only worked with sharded_jit and has
been removed from pmap. Please migrate to pjit and remove global_arg_shapes
from pmap.
Remove Cuda 11.4 support. JAX from 0.4.8 release will support cuda 11.8 and cuda 12 PiperOrigin-RevId: 520056811
2023-03-28 09:53:57 -07:00
Finish jax and jaxlib 0.4.7 release PiperOrigin-RevId: 519839723
2023-03-27 15:06:06 -07:00
## jax 0.4.7 (March 27, 2023)
Update versions and changelog for jax + jaxlib 0.4.6 release
2023-03-09 14:46:11 -08:00
Error if jax_array or jax_jit_pjit_api_merge is set to False. PiperOrigin-RevId: 517485597
2023-03-17 12:57:18 -07:00
* Changes
* As per https://jax.readthedocs.io/en/latest/jax_array_migration.html#jax-array-migration
`jax.config.jax_array` cannot be disabled anymore.
* `jax.config.jax_jit_pjit_api_merge` cannot be disabled anymore.
[jax2tf] Update CHANGELOG for native serialization. PiperOrigin-RevId: 517994283
2023-03-20 09:42:59 -07:00
* {func}`jax.experimental.jax2tf.convert` now supports the `native_serialization`
parameter to use JAX's native lowering to StableHLO to obtain a
StableHLO module for the entire JAX function instead of lowering each JAX
primitive to a TensorFlow op. This simplifies the internals and increases
the confidence that what you serialize matches the JAX native semantics.
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
See [documentation](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/README.md).
[jax2tf] Update CHANGELOG for native serialization. PiperOrigin-RevId: 517994283
2023-03-20 09:42:59 -07:00
As part of this change the config flag `--jax2tf_default_experimental_native_lowering`
has been renamed to `--jax2tf_native_serialization`.
[XLA:Python] Change JAX and the XLA Python extension to get NumPy bfloat16/float8 types from ml_dtypes. PiperOrigin-RevId: 518830467
2023-03-23 05:12:57 -07:00
* JAX now depends on `ml_dtypes`, which contains definitions of NumPy types
like bfloat16. These definitions were previously internal to JAX, but have
been split into a separate package to facilitate sharing them with other
projects.
Increase minimum NumPy version to 1.21. Also increase minimum SciPy version to 1.7, which was released just before NumPy 1.21.
2023-02-06 11:32:28 -05:00
* JAX now requires NumPy 1.21 or newer and SciPy 1.7 or newer.
Error if jax_array or jax_jit_pjit_api_merge is set to False. PiperOrigin-RevId: 517485597
2023-03-17 12:57:18 -07:00
jnp.ndarray.at: deprecate passing additional arguments by position
2023-03-13 10:04:39 -07:00
* Deprecations
Mark jax.numpy.DeviceArray as deprecated. Use jax.Array instead. PiperOrigin-RevId: 516835920
2023-03-15 08:45:09 -07:00
* The type `jax.numpy.DeviceArray` is deprecated. Use `jax.Array` instead,
for which it is an alias.
Mark jax.interpreters.pxla.ShardedDeviceArray as deprecated. PiperOrigin-RevId: 518241326
2023-03-21 05:13:21 -07:00
* The type `jax.interpreters.pxla.ShardedDeviceArray` is deprecated. Use
`jax.Array` instead.
jnp.ndarray.at: deprecate passing additional arguments by position
2023-03-13 10:04:39 -07:00
* Passing additional arguments to {func}`jax.numpy.ndarray.at` by position is deprecated.
For example, instead of `x.at[i].get(True)`, use `x.at[i].get(indices_are_sorted=True)`
[Jax cleanup] * Remove lower_xla_callable and all related functions * Remove pxla.device_put * Remove dispatch.device_put_handlers PiperOrigin-RevId: 517249345
2023-03-16 15:46:57 -07:00
* `jax.interpreters.xla.device_put` is deprecated. Please use `jax.device_put`.
* `jax.interpreters.pxla.device_put` is deprecated. Please use `jax.device_put`.
Add deprecation warning for FROM_GDA usage since that argument is not required anymore. PiperOrigin-RevId: 519781715
2023-03-27 11:32:30 -07:00
* `jax.experimental.pjit.FROM_GDA` is deprecated. Please pass in sharded
jax.Arrays as input and remove the `in_shardings` argument to pjit since
it is optional.
jnp.ndarray.at: deprecate passing additional arguments by position
2023-03-13 10:04:39 -07:00
Finish jax and jaxlib 0.4.7 release PiperOrigin-RevId: 519839723
2023-03-27 15:06:06 -07:00
## jaxlib 0.4.7 (March 27, 2023)
Update versions and changelog for jax + jaxlib 0.4.6 release
2023-03-09 14:46:11 -08:00
[XLA:Python] Change JAX and the XLA Python extension to get NumPy bfloat16/float8 types from ml_dtypes. PiperOrigin-RevId: 518830467
2023-03-23 05:12:57 -07:00
Changes:
* jaxlib now depends on `ml_dtypes`, which contains definitions of NumPy types
like bfloat16. These definitions were previously internal to JAX, but have
been split into a separate package to facilitate sharing them with other
projects.
Update versions and changelog for jax + jaxlib 0.4.6 release
2023-03-09 14:46:11 -08:00
## jax 0.4.6 (Mar 9, 2023)
Bump version and changelog after jax 0.4.5 release
2023-03-02 16:08:34 -08:00
add key path change to changelog
2023-03-03 18:05:37 -08:00
* Changes
fix typo
2023-03-05 13:14:57 -08:00
* `jax.tree_util` now contain a set of APIs that allow user to define keys for their
add key path change to changelog
2023-03-03 18:05:37 -08:00
custom pytree node. This includes:
[shape_poly, call_tf] Some improvements for call_tf in a shape polymorphic program This is another attempt to land a rolled-back change https://github.com/google/jax/pull/14734 (cl/514070997). See b/272154366 for more details. The use case for call_tf with shape polymorphism is when we have a JAX program that calls into TF function, and we want to serialize the JAX program with some shapes unknown. Previously this use case did not work, except in the special case when the output shape of the called TF function returns statically known shapes. The idea is that we allow the user of call_tf to specify the output shape. This can be done even in presence of shape polymorphism, by writing the output shape as an expression in terms of the input shapes. This is what other JAX primitives do, e.g., concat, so we are simply enabling call_tf to get the same behavior. This change should be enough for old-style jax2tf, but will require more work for native serialization. We also removed some old code that was trying to workaround some limitations in shape inference in TF. I think that those workarounds are ugly, and I am prepared to give error messages rather than keep that code. So far no tests fail. PiperOrigin-RevId: 515137407
2023-03-08 14:09:35 -08:00
* `tree_flatten_with_path` that flattens a tree and return not only each leaf but
add key path change to changelog
2023-03-03 18:05:37 -08:00
also their key paths.
Fix Typos in CHANGELOG.md This PR fixes the typos in Change log documentation
2024-03-12 13:57:07 +05:30
* `tree_map_with_path` that can map a function that takes the key path as an argument.
* `register_pytree_with_keys` to register how the key path and leaves should looks
add key path change to changelog
2023-03-03 18:05:37 -08:00
like in a custom pytree node.
* `keystr` that pretty-prints a key path.
[shape_poly, call_tf] Some improvements for call_tf in a shape polymorphic program This is another attempt to land a rolled-back change https://github.com/google/jax/pull/14734 (cl/514070997). See b/272154366 for more details. The use case for call_tf with shape polymorphism is when we have a JAX program that calls into TF function, and we want to serialize the JAX program with some shapes unknown. Previously this use case did not work, except in the special case when the output shape of the called TF function returns statically known shapes. The idea is that we allow the user of call_tf to specify the output shape. This can be done even in presence of shape polymorphism, by writing the output shape as an expression in terms of the input shapes. This is what other JAX primitives do, e.g., concat, so we are simply enabling call_tf to get the same behavior. This change should be enough for old-style jax2tf, but will require more work for native serialization. We also removed some old code that was trying to workaround some limitations in shape inference in TF. I think that those workarounds are ugly, and I am prepared to give error messages rather than keep that code. So far no tests fail. PiperOrigin-RevId: 515137407
2023-03-08 14:09:35 -08:00
* {func}`jax2tf.call_tf` has a new parameter `output_shape_dtype` (default `None`)
that can be used to declare the output shape and type of the result. This enables
{func}`jax2tf.call_tf` to work in the presence of shape polymorphism. ({jax-issue}`#14734`).
add key path change to changelog
2023-03-03 18:05:37 -08:00
* Deprecations
* The old key-path APIs in `jax.tree_util` are deprecated and will be removed 3 months
from Mar 10 2023:
fix typo
2023-03-05 13:14:57 -08:00
* `register_keypaths`: use {func}`jax.tree_util.register_pytree_with_keys` instead.
add key path change to changelog
2023-03-03 18:05:37 -08:00
* `AttributeKeyPathEntry` : use `GetAttrKey` instead.
* `GetitemKeyPathEntry` : use `SequenceKey` or `DictKey` instead.
Update versions and changelog for jax + jaxlib 0.4.6 release
2023-03-09 14:46:11 -08:00
## jaxlib 0.4.6 (Mar 9, 2023)
Bump version and changelog after jax 0.4.5 release
2023-03-02 16:08:34 -08:00
## jax 0.4.5 (Mar 2, 2023)
Finish jax and jaxlib 0.4.4 release PiperOrigin-RevId: 510234171
2023-02-16 13:48:13 -08:00
Rename `jax.sharding.OpShardingSharding` to `jax.sharding.GSPMDSharding`. `jax.sharding.OpShardingSharding` will be removed in 3 months from Feb 17, 2023. PiperOrigin-RevId: 510556189
2023-02-17 17:10:27 -08:00
* Deprecations
* `jax.sharding.OpShardingSharding` has been renamed to `jax.sharding.GSPMDSharding`.
`jax.sharding.OpShardingSharding` will be removed in 3 months from Feb 17, 2023.
Deprecate three jax.Array methods: - jax.Array.broadcast: use lax.broadcast instead - jax.Array.broadcast_in_dim: use lax.broadcast_in_dim instead - jax.Array.split: use jnp.split instead These are removed because they are not part of the np.ndarray API.
2023-02-23 16:15:09 -08:00
* The following `jax.Array` methods are deprecated and will be removed 3 months from
Feb 23 2023:
* `jax.Array.broadcast`: use {func}`jax.lax.broadcast` instead.
* `jax.Array.broadcast_in_dim`: use {func}`jax.lax.broadcast_in_dim` instead.
* `jax.Array.split`: use {func}`jax.numpy.split` instead.
Rename `jax.sharding.OpShardingSharding` to `jax.sharding.GSPMDSharding`. `jax.sharding.OpShardingSharding` will be removed in 3 months from Feb 17, 2023. PiperOrigin-RevId: 510556189
2023-02-17 17:10:27 -08:00
Finish jax and jaxlib 0.4.4 release PiperOrigin-RevId: 510234171
2023-02-16 13:48:13 -08:00
## jax 0.4.4 (Feb 16, 2023)
Post 0.4.3 release updates
2023-02-08 10:08:57 -08:00
Default `jax_jit_pjit_api_merge` to True. This means that the implementation of jit and pjit have been merged but they still remain separate APIs due to the semantic difference of how they behave under the `Mesh` context manager. This changes the internals of JAX without affecting any public API. Before, `jit` was a final style primitive. This means that the creation of jaxpr was delayed as much as possible and transformations were stacked on top of each other. With the `jit`-`pjit` implementation merge, `jit` becomes an initial style primitive which means that we trace to jaxpr as early as possible. For more information see [this section in autodidax](https://jax.readthedocs.io/en/latest/autodidax.html#on-the-fly-final-style-and-staged-initial-style-processing). Moving to initial style should simplify JAX's internals and make development of features like dynamic shapes, etc easier. PiperOrigin-RevId: 508143501
2023-02-08 11:55:10 -08:00
* Changes
* The implementation of `jit` and `pjit` has been merged. Merging jit and pjit
changes the internals of JAX without affecting the public API of JAX.
Before, `jit` was a final style primitive. Final style means that the creation
of jaxpr was delayed as much as possible and transformations were stacked
on top of each other. With the `jit`-`pjit` implementation merge, `jit`
becomes an initial style primitive which means that we trace to jaxpr
as early as possible. For more information see
[this section in autodidax](https://jax.readthedocs.io/en/latest/autodidax.html#on-the-fly-final-style-and-staged-initial-style-processing).
Moving to initial style should simplify JAX's internals and make
development of features like dynamic shapes, etc easier.
You can disable it only via the environment variable i.e.
`os.environ['JAX_JIT_PJIT_API_MERGE'] = '0'`.
The merge must be disabled via an environment variable since it affects JAX
at import time so it needs to be disabled before jax is imported.
Change the `axis_resources` argument of `with_sharding_constraint` to `shardings` to match `pjit` and `jit`. PiperOrigin-RevId: 509275107
2023-02-13 10:53:21 -08:00
* `axis_resources` argument of `with_sharding_constraint` is deprecated.
Please use `shardings` instead. There is no change needed if you were using
`axis_resources` as an arg. If you were using it as a kwarg, then please
use `shardings` instead. `axis_resources` will be removed after 3 months
from Feb 13, 2023.
Expose jax.typing & update docs
2023-02-13 15:53:08 -08:00
* added the {mod}`jax.typing` module, with tools for type annotations of JAX
functions.
Deprecate uses of jax.experimental.pjit.NamedSharding and jax.experimental.pjit.PartitionSpec. Use the aliases under jax.sharding instead. PiperOrigin-RevId: 509837529
2023-02-15 08:13:58 -08:00
* The following names have been deprecated:
* `jax.xla.Device` and `jax.interpreters.xla.Device`: use `jax.Device`.
[JAX] Deprecate jax.experimental.maps.Mesh. PiperOrigin-RevId: 509852142
2023-02-15 09:15:11 -08:00
* `jax.experimental.maps.Mesh`. Use `jax.sharding.Mesh`
instead.
Deprecate uses of jax.experimental.pjit.NamedSharding and jax.experimental.pjit.PartitionSpec. Use the aliases under jax.sharding instead. PiperOrigin-RevId: 509837529
2023-02-15 08:13:58 -08:00
* `jax.experimental.pjit.NamedSharding`: use `jax.sharding.NamedSharding`.
* `jax.experimental.pjit.PartitionSpec`: use `jax.sharding.PartitionSpec`.
* `jax.interpreters.pxla.Mesh`: use `jax.sharding.Mesh`.
* `jax.interpreters.pxla.PartitionSpec`: use `jax.sharding.PartitionSpec`.
jax.numpy reductions: require initial to be a scalar This follows the requirements of numpy's reduction API. Non-scalar initial values can be implemented via .
2023-02-14 15:36:18 -08:00
* Breaking Changes
* the `initial` argument to reduction functions like :func:`jax.numpy.sum`
is now required to be a scalar, consistent with the corresponding NumPy API.
Fix Typos in CHANGELOG.md This PR fixes the typos in Change log documentation
2024-03-12 13:57:07 +05:30
The previous behavior of broadcasting the output against non-scalar `initial`
jax.numpy reductions: require initial to be a scalar This follows the requirements of numpy's reduction API. Non-scalar initial values can be implemented via .
2023-02-14 15:36:18 -08:00
values was an unintentional implementation detail ({jax-issue}`#14446`).
Default `jax_jit_pjit_api_merge` to True. This means that the implementation of jit and pjit have been merged but they still remain separate APIs due to the semantic difference of how they behave under the `Mesh` context manager. This changes the internals of JAX without affecting any public API. Before, `jit` was a final style primitive. This means that the creation of jaxpr was delayed as much as possible and transformations were stacked on top of each other. With the `jit`-`pjit` implementation merge, `jit` becomes an initial style primitive which means that we trace to jaxpr as early as possible. For more information see [this section in autodidax](https://jax.readthedocs.io/en/latest/autodidax.html#on-the-fly-final-style-and-staged-initial-style-processing). Moving to initial style should simplify JAX's internals and make development of features like dynamic shapes, etc easier. PiperOrigin-RevId: 508143501
2023-02-08 11:55:10 -08:00
Finish jax and jaxlib 0.4.4 release PiperOrigin-RevId: 510234171
2023-02-16 13:48:13 -08:00
## jaxlib 0.4.4 (Feb 16, 2023)
Drop support for NVIDIA Kepler series GPUs in jaxlib builds.
2023-02-10 14:15:15 -05:00
* Breaking changes
* Support for NVIDIA Kepler series GPUs has been removed from the default
`jaxlib` builds. If Kepler support is needed, it is still possible to
build `jaxlib` from source with Kepler support (via the
`--cuda_compute_capabilities=sm_35` option to `build.py`), however note
that CUDA 12 has completely dropped support for Kepler GPUs.
Post 0.4.3 release updates
2023-02-08 10:08:57 -08:00
## jax 0.4.3 (Feb 8, 2023)
Removed deprecated polar_unitary as per comment.
2023-01-26 23:01:06 +00:00
* Breaking changes
Post 0.4.3 release updates
2023-02-08 10:08:57 -08:00
* Deleted {func}`jax.scipy.linalg.polar_unitary`, which was a deprecated JAX
Removed deprecated polar_unitary as per comment.
2023-01-26 23:01:06 +00:00
extension to the scipy API. Use {func}`jax.scipy.linalg.polar` instead.
Updates for jax + jaxlib 0.4.2 release
2023-01-20 19:04:46 +00:00
Added scipy.stats.rankdata
2023-02-07 12:07:00 -05:00
* Changes
* Added {func}`jax.scipy.stats.rankdata`.
Post 0.4.3 release updates
2023-02-08 10:08:57 -08:00
## jaxlib 0.4.3 (Feb 8, 2023)
DOC: add is_ready() to CHANGELOG
2023-02-23 11:56:48 -08:00
* `jax.Array` now has the non-blocking `is_ready()` method, which returns `True`
if the array is ready (see also {func}`jax.block_until_ready`).
Updates for jax + jaxlib 0.4.2 release
2023-01-20 19:04:46 +00:00
Update versions for jaxlib 0.4.2 release. I also screwed up the CHANGELOG before (I shouldn't have added a date), so I'm fixing the dates now.
2023-01-26 01:09:26 +00:00
## jax 0.4.2 (Jan 24, 2023)
Finish jax and jaxlib 0.4.2 release PiperOrigin-RevId: 495068000
2022-12-13 10:50:25 -08:00
Delete `jax.experimental.callback` PiperOrigin-RevId: 501760507
2023-01-12 22:57:49 -08:00
* Breaking changes
* Deleted `jax.experimental.callback`
[shape_poly] Generalize binary operations with symbolic dimensions. Previously binary operations involving symbolic dimensions would work only when the other operand is convertible to a symbolic dimension, e.g., an integer. This resulted in errors when trying "x.shape[0] * 3.5" and the recourse was to ask the user to add an explicit "jnp.array(x.shape[0])". Now we allow binary operations with any operand and the "jnp.array" is added automatically if the other operand is not an integer or a symbolic dimension. This means that instead of an error they may be an error downstream if one tries to use the result as a dimension. There is one known case where JAX works with static shapes and with the previous behavior, but will fail now. When you operate on `np.ndarray` and symbolic dimension, previously this was kept as a `np.ndarray` but not it is turned into a JAX array. The following program will now fail if `x.shape[0]` is a symbolic dimension.: `jnp.ones(np.arange(5) * x.shape[0])` Instead you should write `jnp.ones([i * x.shape[0] for i in range(5)])`
2023-01-18 12:27:02 +02:00
* 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`).
Raise an error when attempting to mutate Jaxpr objects
2023-01-23 09:37:58 -08:00
* jaxpr objects now raise an error on attribute setting in order to avoid
problematic mutations ({jax-issue}`14102`)
Delete `jax.experimental.callback` PiperOrigin-RevId: 501760507
2023-01-12 22:57:49 -08:00
[call_tf] Add has_side_effects parameter The CallTfEffect was added recently as an internal workaround for DCE removing instances of call_tf. Here we add a parameter to `call_tf` to be able to declare if the called computation is effectful and should not be removed by DCE.
2023-01-12 08:44:53 +01:00
* 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`).
[shape_poly] Add better support for division, and working with strides Previously, division was only supported in certain situation, and this led to errors, e.g., when using strides. Now we generalize the polynomials to also include "floordiv(E, E)" and "mod(E, E)" as atoms, in addition to dimension variables. A symbolic dimension is now a sum of products of atoms. (We also changed the documentation to use symbolic dimension instead of dimension polynomials).
2023-01-18 12:27:02 +02:00
* 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`).
[call_tf] Add has_side_effects parameter The CallTfEffect was added recently as an internal workaround for DCE removing instances of call_tf. Here we add a parameter to `call_tf` to be able to declare if the called computation is effectful and should not be removed by DCE.
2023-01-12 08:44:53 +01:00
Update versions for jaxlib 0.4.2 release. I also screwed up the CHANGELOG before (I shouldn't have added a date), so I'm fixing the dates now.
2023-01-26 01:09:26 +00:00
## jaxlib 0.4.2 (Jan 24, 2023)
Updates for jax + jaxlib 0.4.2 release
2023-01-20 19:04:46 +00:00
* Changes
* Set JAX_USE_PJRT_C_API_ON_TPU=1 to enable new Cloud TPU runtime, featuring
automatic device memory defragmentation.
Finish jax and jaxlib 0.4.2 release PiperOrigin-RevId: 495068000
2022-12-13 10:50:25 -08:00
## jax 0.4.1 (Dec 13, 2022)
Finish jax and jaxlib release 0.4.0 PiperOrigin-RevId: 494833878
2022-12-12 14:42:48 -08:00
Add the `jax.Array` change to the changelog. PiperOrigin-RevId: 488929264
2022-11-16 06:55:34 -08:00
* Changes
Drop support for Python 3.7
2022-11-29 15:01:47 -08:00
* Support for Python 3.7 has been dropped, in accordance with JAX's
{ref}`version-support-policy`.
Add the `jax.Array` change to the changelog. PiperOrigin-RevId: 488929264
2022-11-16 06:55:34 -08:00
* 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,
[docs] Replace one more jax_Array.html reference I missed this in #13479, thanks @yashk2810 for flagging!
2022-12-01 20:50:55 +00:00
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](https://jax.readthedocs.io/en/latest/jax_array_migration.html) can
help you migrate your codebase to `jax.Array`. You can also look at the
in jax.Array notebook, polish beginning and tweak title and some wording
2022-12-10 19:59:14 -08:00
[Distributed arrays and automatic parallelization](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html)
Add the `jax.Array` change to the changelog. PiperOrigin-RevId: 488929264
2022-11-16 06:55:34 -08:00
tutorial to understand the new concepts.
Move `PartitionSpec` and `Mesh` out of experimental and into the `sharding` namespace. The new API endpoint is `jax.sharding.PartitionSpec` and `jax.sharding.Mesh`. PiperOrigin-RevId: 492358238
2022-12-01 19:28:02 -08:00
* `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.
Move `with_sharding_constraint` out of experimental into `jax.lax` namespace. PiperOrigin-RevId: 494635809
2022-12-11 22:54:39 -08:00
* `with_sharding_constraint`s new public endpoint is
`jax.lax.with_sharding_constraint`.
Do not mirror JAX config options back to ABSL flags. Currently when JAX config values are configured via ABSL, we use the ABSL flags as a source of truth: if we read or write the JAX config option, we read or write the corresponding ABSL flag. This works but has the unfortunate downside that ABSL flags are relatively slow to read, which slows down JAX every time we read a configuration option. However, there's fundamentally no reason we are mirroring the JAX configuration options back to ABSL in the first place. We can use ABSL flag parsing as a way only to populate the JAX configuration values. The downside is that if someone changes the ABSL flag values after parsing, that change will not be reflected in JAX's config values. JAX config changes after ABSL flags have been parsed must be made via the `jax.config.update()` API. This gives a decent improvement on the device_put benchmark: ``` name old cpu/op new cpu/op delta device_put 79.5µs ± 6% 69.4µs ± 7% -12.73% (p=0.000 n=10+9) name old time/op new time/op delta device_put 79.5µs ± 6% 69.4µs ± 7% -12.73% (p=0.000 n=10+9) ``` PiperOrigin-RevId: 492519085
2022-12-02 11:36:51 -08:00
* 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.
Mark arguments to ufuncs as positional-only. PiperOrigin-RevId: 493311821
2022-12-06 08:23:40 -08:00
* 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.
Deprecate jnp.msort following deprecation of numpy.msort
2022-12-07 10:08:18 -08:00
* `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)`.
Finish jax and jaxlib release 0.3.25. The next release will be 0.4.0 (since jax.Array will be enabled in that release) PiperOrigin-RevId: 488672395
2022-11-15 09:02:23 -08:00
Finish jax and jaxlib 0.4.2 release PiperOrigin-RevId: 495068000
2022-12-13 10:50:25 -08:00
## jaxlib 0.4.1 (Dec 13, 2022)
Update CHANGELOG to indicate that 0.4.0 was yanked. PiperOrigin-RevId: 494868884
2022-12-12 17:01:54 -08:00
updated jaxlib CHANGELOG
2022-11-28 10:37:42 -08:00
* Changes
Drop support for Python 3.7
2022-11-29 15:01:47 -08:00
* Support for Python 3.7 has been dropped, in accordance with JAX's
{ref}`version-support-policy`.
updated jaxlib CHANGELOG
2022-11-28 10:37:42 -08:00
* The behavior of `XLA_PYTHON_CLIENT_MEM_FRACTION=.XX` has been changed to allocate XX% of
the total GPU memory instead of the previous behavior of using currently available GPU memory
to calculate preallocation. Please refer to
[GPU memory allocation](https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html) for
more details.
Remove long-deprecated method .block_host_until_ready(). PiperOrigin-RevId: 492571809
2022-12-02 15:13:55 -08:00
* The deprecated method `.block_host_until_ready()` has been removed. Use
`.block_until_ready()` instead.
updated jaxlib CHANGELOG
2022-11-28 10:37:42 -08:00
Update CHANGELOG to indicate that 0.4.0 was yanked. PiperOrigin-RevId: 494868884
2022-12-12 17:01:54 -08:00
## jax 0.4.0 (Dec 12, 2022)
* The release was yanked.
## jaxlib 0.4.0 (Dec 12, 2022)
* The release was yanked.
Finish jax and jaxlib release 0.3.25. The next release will be 0.4.0 (since jax.Array will be enabled in that release) PiperOrigin-RevId: 488672395
2022-11-15 09:02:23 -08:00
## jax 0.3.25 (Nov 15, 2022)
Add support for the hermitian option on jnp.linalg.pinv. Improve the pinv implementation to avoid computing an unnecessary reduction: svd sorts its singular values so we don't need to use amax() to find the largest one. Avoid explicitly forming the identity matrix in the pinv JVP.
2022-11-07 11:39:19 -05:00
* Changes
* {func}`jax.numpy.linalg.pinv` now supports the `hermitian` option.
Add support for Hessenberg and tridiagonal matrix reductions on CPU. * Implement jax.scipy.linalg.hessenberg and jax.lax.linalg.hessenberg. * Export what was previously jax._src.lax.linalg.orgqr as jax.lax.linalg.householder_product, since it can be used with some minor tweaks to compute the unitary matrix of a Hessenberg reduction. * Implement jax.lax.linalg.tridiagonal, which is the symmetric (Hermitian) equivalent of Hessenberg reduction. None of these primitives are differentiable at the moment. PiperOrigin-RevId: 487224934
2022-11-09 06:23:22 -08:00
* {func}`jax.scipy.linalg.hessenberg` is now supported on CPU only. Requires
jaxlib > 0.3.24.
* New functions {func}`jax.lax.linalg.hessenberg`,
{func}`jax.lax.linalg.tridiagonal`, and
Add several recent changes to the CHANGELOG. PiperOrigin-RevId: 488362198
2022-11-14 07:38:39 -08:00
{func}`jax.lax.linalg.householder_product` were added. Householder reduction
is currently CPU-only and tridiagonal reductions are supported on CPU and
GPU only.
* The gradients of `svd` and `jax.numpy.linalg.pinv` are now computed more
economically for non-square matrices.
Update changelog w/ info about deleting `jax_experimental_name_stack`
2022-11-11 14:00:33 -08:00
* Breaking Changes
* Deleted the `jax_experimental_name_stack` config option.
Convert string axis name into tuple of strings in Mesh constructor PiperOrigin-RevId: 487930412
2022-11-11 15:23:44 -08:00
* Convert a string `axis_names` arguments to the
{class}`jax.experimental.maps.Mesh` constructor into a singleton tuple
instead of unpacking the string into a sequence of character axis names.
Finish the release of jax and jaxlib 0.3.24 PiperOrigin-RevId: 486162090
2022-11-04 09:41:29 -07:00
Finish jax and jaxlib release 0.3.25. The next release will be 0.4.0 (since jax.Array will be enabled in that release) PiperOrigin-RevId: 488672395
2022-11-15 09:02:23 -08:00
## jaxlib 0.3.25 (Nov 15, 2022)
Add support for Hessenberg and tridiagonal matrix reductions on CPU. * Implement jax.scipy.linalg.hessenberg and jax.lax.linalg.hessenberg. * Export what was previously jax._src.lax.linalg.orgqr as jax.lax.linalg.householder_product, since it can be used with some minor tweaks to compute the unitary matrix of a Hessenberg reduction. * Implement jax.lax.linalg.tridiagonal, which is the symmetric (Hermitian) equivalent of Hessenberg reduction. None of these primitives are differentiable at the moment. PiperOrigin-RevId: 487224934
2022-11-09 06:23:22 -08:00
* Changes
Add several recent changes to the CHANGELOG. PiperOrigin-RevId: 488362198
2022-11-14 07:38:39 -08:00
* Added support for tridiagonal reductions on CPU and GPU.
* Added support for upper Hessenberg reductions on CPU.
* Bugs
* Fixed a bug that meant that frames in tracebacks captured by JAX were
incorrectly mapped to source lines under Python 3.10+
Finish the release of jax and jaxlib 0.3.24 PiperOrigin-RevId: 486162090
2022-11-04 09:41:29 -07:00
## jax 0.3.24 (Nov 4, 2022)
[JAX] Change the default pmap() ordering to match the ordering of jax.devices() for single-process TPU jobs. PiperOrigin-RevId: 484062717
2022-10-26 13:55:32 -07:00
* Changes
Avoid loading scipy eagerly. scipy accounts for around 400ms of the 900ms of JAX's import time. By loading scipy lazily, we can improve the timing of `import jax` down to about 500ms.
2022-10-06 01:43:54 +00:00
* JAX should be faster to import. We now import scipy lazily, which accounted
for a significant fraction of JAX's import time.
Add new config `jax_persistent_cache_min_compile_time_secs`. This replaces `jax_persistent_cache_min_instruction_count` introduced in https://github.com/google/jax/pull/12798, since gating on the compile time seems strictly better than gating on the instruction count (except maybe that the instruction count is more deterministic, but I don't think that's a big deal). I defaulted to 1 second as the minimum threshold based on the same flax wmt example (https://github.com/google/flax/tree/main/examples/wmt) numbers from name | instruction_count | compile_time_secs ---- | ----------------- | ----------------- `broadcast_in_dim` | 2 | 0.01633763313 `convert_element_type` | 2 | 0.01704716682 `reshape` | 2 | 0.01730203629 `_squareit` | 2 | 0.01730823517 `broadcast_in_dim` | 2 | 0.0182030201 `convert_element_type` | 2 | 0.01982188225 `concatenate` | 2 | 0.02102327347 `true_divide` | 2 | 0.02172231674 `broadcast_in_dim` | 2 | 0.02370619774 `broadcast_in_dim` | 2 | 0.02393102646 `broadcast_in_dim` | 2 | 0.02488565445 `broadcast_in_dim` | 2 | 0.03395628929 `broadcast_in_dim` | 2 | 0.03428125381 `broadcast_in_dim` | 2 | 0.0394551754 `shift_right_logical` | 2 | 0.06500506401 `<lambda>` | 3 | 0.01793265343 `_unstack` | 5 | 0.01975226402 `_reduce_sum` | 5 | 0.0210878849 `_reduce_sum` | 5 | 0.02416801453 `_multi_slice` | 9 | 0.09065580368 `_threefry_split` | 232 | 0.09037566185 `_threefry_split` | 232 | 0.09161829948 `<unnamed wrapped function>` | 2668 | 7.701903343 `<unnamed wrapped function>` | 3455 | 17.57672167 `<unnamed wrapped function>` | 46580 | 166.2570884 `init` | 60361 | 26.35722399 `<unnamed wrapped function>` | 78010 | 3.879326344 Also adds new float config functionality.
2022-10-28 23:53:30 +00:00
* Setting the env var `JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS=$N` can be
used to limit the number of cache entries written to the persistent cache.
By default, computations that take 1 second or more to compile will be
cached.
Add docs & changelog for jax.scipy.stats.mode
2022-10-20 15:55:57 -07:00
* Added {func}`jax.scipy.stats.mode`.
[JAX] Change the default pmap() ordering to match the ordering of jax.devices() for single-process TPU jobs. PiperOrigin-RevId: 484062717
2022-10-26 13:55:32 -07:00
* The default device order used by `pmap` on TPU if no order is specified now
matches `jax.devices()` for single-process jobs. Previously the
two orderings differed, which could lead to unnecessary copies or
out-of-memory errors. Requiring the orderings to agree simplifies matters.
* Breaking Changes
jnp.gradient: call check_arraylike on inputs & clean-up implementation
2022-10-24 15:27:33 -07:00
* {func}`jax.numpy.gradient` now behaves like most other functions in {mod}`jax.numpy`,
and forbids passing lists or tuples in place of arrays ({jax-issue}`#12958`)
Call _check_arraylike for jnp.linalg & jnp.fft functions
2022-10-27 09:01:28 -07:00
* Functions in {mod}`jax.numpy.linalg` and {mod}`jax.numpy.fft` now uniformly
require inputs to be array-like: i.e. lists and tuples cannot be used in place
of arrays. Part of {jax-issue}`#7737`.
Rename `ReshapeableDevicesSharding` to `PositionalSharding` and add an alias `NamedSharding` for `MeshPspecSharding`. `MeshPspecSharding` name will be replaced with `NamedSharding` in 3 months. PiperOrigin-RevId: 485753078
2022-11-02 19:12:32 -07:00
* Deprecations
* `jax.sharding.MeshPspecSharding` has been renamed to `jax.sharding.NamedSharding`.
`jax.sharding.MeshPspecSharding` name will be removed in 3 months.
Update version.py and CHANGELOG for jax 0.3.23 release
2022-10-12 11:09:15 -07:00
Finish the release of jax and jaxlib 0.3.24 PiperOrigin-RevId: 486162090
2022-11-04 09:41:29 -07:00
## jaxlib 0.3.24 (Nov 4, 2022)
[JAX:CPU] Enable buffer donation on CPU. Fix a bug in PJRT where if a buffer was not owned (e.g., it aliased a NumPy buffer) it could still be donated and that would lead to a use after free. PiperOrigin-RevId: 484001545
2022-10-26 10:07:32 -07:00
* Changes
* Buffer donation now works on CPU. This may break code that marked buffers
for donation on CPU but relied on donation not being implemented.
Update version.py and CHANGELOG for jax 0.3.23 release
2022-10-12 11:09:15 -07:00
## jax 0.3.23 (Oct 12, 2022)
Update Colab TPU driver version
2022-10-11 17:49:10 -07:00
* Changes
* Update Colab TPU driver version for new jaxlib release.
Update version.py, setup.py and CHANGELOG post jax 0.3.22 release
2022-10-11 21:37:53 +00:00
## jax 0.3.22 (Oct 11, 2022)
Update message and change log
2022-10-05 01:01:31 +00:00
* Changes
Address comments for change log
2022-10-05 18:16:49 +00:00
* Add `JAX_PLATFORMS=tpu,cpu` as default setting in TPU initialization,
so JAX will raise an error if TPU cannot be initialized instead of falling
back to CPU. Set `JAX_PLATFORMS=''` to override this behavior and automatically
choose an available backend (the original default), or set `JAX_PLATFORMS=cpu`
to always use CPU regardless of if the TPU is available.
Remove deprecated functionality from jax.test_util PiperOrigin-RevId: 480360504
2022-10-11 08:15:47 -07:00
* Deprecations
* Several test utilities deprecated in JAX v0.3.8 are now removed from
{mod}`jax.test_util`.
Update version and CHANGELOG for jax 0.3.21 release
2022-10-03 22:06:19 +00:00
Update version.py, setup.py and CHANGELOG post jax 0.3.22 release
2022-10-11 21:37:53 +00:00
## jaxlib 0.3.22 (Oct 11, 2022)
Update version and CHANGELOG for jax 0.3.21 release
2022-10-03 22:06:19 +00:00
## jax 0.3.21 (Sep 30, 2022)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.3.20...jax-v0.3.21).
Make persistent compilation cache warn instead of raise an error on cache read/write failures Fixes #12582. Setting the env var `JAX_RAISE_PERSISTENT_CACHE_ERRORS=true` will revert to the original behavior of raising exception instead of warning. Also makes JAX_DUMP_IR_TO work when the persistent cache is enabled.
2022-09-27 20:59:08 +00:00
* Changes
* The persistent compilation cache will now warn instead of raising an
exception on error ({jax-issue}`#12582`), so program execution can continue
if something goes wrong with the cache. Set
`JAX_RAISE_PERSISTENT_CACHE_ERRORS=true` to revert this behavior.
Update version numbers after release.
2022-09-28 18:49:22 +00:00
## jax 0.3.20 (Sep 28, 2022)
jax/jaxlib 0.3.20 release candidate.
2022-09-28 13:33:52 +00:00
* Bug fixes:
* Adds missing `.pyi` files that were missing from the previous release ({jax-issue}`#12536`).
* Fixes an incompatibility between `jax` 0.3.19 and the libtpu version it pinned ({jax-issue}`#12550`). Requires jaxlib 0.3.20.
* Fix incorrect `pip` url in `setup.py` comment ({jax-issue}`#12528`).
Update version and changelog for jax 0.3.18 release
2022-09-26 12:38:32 -07:00
Update version numbers after release.
2022-09-28 18:49:22 +00:00
## jaxlib 0.3.20 (Sep 28, 2022)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jaxlib-v0.3.15...jaxlib-v0.3.20).
jax/jaxlib 0.3.20 release candidate.
2022-09-28 13:33:52 +00:00
* Bug fixes
* Fixes support for limiting the visible CUDA devices via
`jax_cuda_visible_devices` in distributed jobs. This functionality is needed for
the JAX/SLURM integration on GPU ({jax-issue}`#12533`).
Update version and changelog for jax 0.3.19 release
2022-09-27 11:00:27 -07:00
## jax 0.3.19 (Sep 27, 2022)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.3.18...jax-v0.3.19).
jax/jaxlib 0.3.20 release candidate.
2022-09-28 13:33:52 +00:00
* Fixes required jaxlib version.
Update version and changelog for jax 0.3.18 release
2022-09-26 12:38:32 -07:00
Jax and jaxlib 0.3.18 release candidate.
2022-09-26 14:10:57 +00:00
## jax 0.3.18 (Sep 26, 2022)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.3.17...jax-v0.3.18).
mention AOT readiness in changelog
2022-09-02 08:23:57 -07:00
* Changes
* Ahead-of-time lowering and compilation functionality (tracked in
{jax-issue}`#7733`) is stable and public. See [the
overview](https://jax.readthedocs.io/en/latest/aot.html) and the API docs
for {mod}`jax.stages`.
Add initial jax.Array base class for instance checks & annotation
2022-09-23 09:59:46 -07:00
* Introduced {class}`jax.Array`, intended to be used for both `isinstance` checks
and type annotations for array types in JAX. Notice that this included some subtle
changes to how `isinstance` works for {class}`jax.numpy.ndarray` for jax-internal
objects, as {class}`jax.numpy.ndarray` is now a simple alias of {class}`jax.Array`.
Remove jax._src from JAX namespace. This is a JAX-internal name and not subject to any deprecation policy. Please avoid the use of JAX-internal functions outside JAX. PiperOrigin-RevId: 473243243
2022-09-09 07:05:30 -07:00
* Breaking changes
Fix Typos in CHANGELOG.md This PR fixes the typos in Change log documentation
2024-03-12 13:57:07 +05:30
* `jax._src` is no longer imported into the public `jax` namespace.
Remove jax._src from JAX namespace. This is a JAX-internal name and not subject to any deprecation policy. Please avoid the use of JAX-internal functions outside JAX. PiperOrigin-RevId: 473243243
2022-09-09 07:05:30 -07:00
This may break users that were using JAX internals.
Delete soft_pmap as it has no users. Please use `pjit` or `xmap` if you do want soft_pmap. `jax.soft_pmap` is undocumented. If it were documented, a deprecation period would have been provided. PiperOrigin-RevId: 474145090
2022-09-13 15:51:30 -07:00
* `jax.soft_pmap` has been deleted. Please use `pjit` or `xmap` instead.
`jax.soft_pmap` is undocumented. If it were documented, a deprecation period
would have been provided.
update version and changelog for pypi
2022-08-31 11:37:09 -07:00
## jax 0.3.17 (Aug 31, 2022)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.3.16...jax-v0.3.17).
Fix JVP rule for lax.pow()
2022-08-22 11:53:02 -07:00
* Bugs
* Fix corner case issue in gradient of `lax.pow` with an exponent of zero
({jax-issue}`12041`)
delete old remat implementation moved lowering rule logic from remat_impl.py (now deleted) to ad_checkpoint.py
2022-08-16 12:07:27 -07:00
* Breaking changes
* {func}`jax.checkpoint`, also known as {func}`jax.remat`, no longer supports
the `concrete` option, following the previous version's deprecation; see
[JEP 11830](https://jax.readthedocs.io/en/latest/jep/11830-new-remat-checkpoint.html).
Expose pure callback and enable rank polymorphic callbacks
2022-08-15 17:05:27 -07:00
* Changes
Apply suggestions from code review Co-authored-by: Matthew Johnson <mattjj@google.com>
2022-08-17 10:43:50 -07:00
* Added {func}`jax.pure_callback` that enables calling back to pure Python functions from compiled functions (e.g. functions decorated with `jax.jit` or `jax.pmap`).
Remove deprecated array.tile() method
2022-08-16 11:35:13 -07:00
* Deprecations:
* The deprecated `DeviceArray.tile()` method has been removed. Use {func}`jax.numpy.tile`
({jax-issue}`#11944`).
[JAX] Deprecate .to_py() property on arrays. Implement __array__ instead. .to_py() was something of an accidental export from the JAX array classes. There are other mechanisms to turn a JAX array into a NumPy array, including `np.asarray(x)` and `jax.device_get(x)`. Deprecate this mechanism because it is redundant. PiperOrigin-RevId: 469984029
2022-08-25 07:27:54 -07:00
* `DeviceArray.to_py()` has been deprecated. Use `np.asarray(x)` instead.
Update version + changelog
2022-08-11 17:10:21 -07:00
## jax 0.3.16
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.3.15...main).
Increase minimum NumPy version to 1.20. Per NEP 29, support for 1.19 ended on Jun 21, 2022.
2022-08-06 14:49:09 +00:00
* Breaking changes
* Support for NumPy 1.19 has been dropped, per the
[deprecation policy](https://jax.readthedocs.io/en/latest/deprecation.html).
Please upgrade to NumPy 1.20 or newer.
Update changelog for jax/jaxlib v0.3.15 release
2022-07-25 09:47:44 -07:00
* Changes
Update CHANGELOG.md
2022-08-02 22:28:52 -07:00
* Added {mod}`jax.debug` that includes utilities for runtime value debugging such at {func}`jax.debug.print` and {func}`jax.debug.breakpoint`.
* Added new documentation for [runtime value debugging](debugging/index)
Update changelog for jax/jaxlib v0.3.15 release
2022-07-25 09:47:44 -07:00
* Deprecations
* {func}`jax.mask` {func}`jax.shapecheck` APIs have been removed.
See {jax-issue}`#11557`.
Remove jax.experimental.loops. See CHANGELOG PiperOrigin-RevId: 463297399
2022-07-26 03:39:17 -07:00
* {mod}`jax.experimental.loops` has been removed. See {jax-issue}`#10278`
set new jax.remat / jax.checkpoint to be on-by-default
2022-08-10 08:26:47 -07:00
for an alternative API.
Remove deprecated function jax.tree_util.tree_multimap
2022-07-06 09:06:16 -07:00
* {func}`jax.tree_util.tree_multimap` has been removed. It has been deprecated since
JAX release 0.3.5, and {func}`jax.tree_util.tree_map` is a direct replacement.
Remove deprecated jax.experimental.stax The new location is jax.example_libraries.stax
2022-08-01 14:58:41 -07:00
* Removed `jax.experimental.stax`; it has long been a deprecated alias of
{mod}`jax.example_libraries.stax`.
Remove deprecated jax.experimental.optimizers The new location is jax.example_libraries.optimizers
2022-08-02 14:41:50 -07:00
* Removed `jax.experimental.optimizers`; it has long been a deprecated alias of
{mod}`jax.example_libraries.optimizers`.
set new jax.remat / jax.checkpoint to be on-by-default
2022-08-10 08:26:47 -07:00
* {func}`jax.checkpoint`, also known as {func}`jax.remat`, has a new
implementation switched on by default, meaning the old implementation is
deprecated; see [JEP 11830](https://jax.readthedocs.io/en/latest/jep/11830-new-remat-checkpoint.html).
Update changelog for jax/jaxlib v0.3.15 release
2022-07-25 09:47:44 -07:00
## jax 0.3.15 (July 22, 2022)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.3.14...jax-v0.3.15).
jnp.roots: better support for computation under JIT
2022-06-23 14:48:53 -07:00
* Changes
jax.test_util: remove deprecated test classes. JaxTestCase and JaxTestLoader were deprecated in jax v0.3.1, released Feb 2022.
2022-06-27 11:04:50 -07:00
* `JaxTestCase` and `JaxTestLoader` have been removed from `jax.test_util`. These
classes have been deprecated since v0.3.1 ({jax-issue}`#11248`).
Implementation of jax.scipy.stats.gaussian_kde
2022-06-23 10:00:54 -04:00
* Added {class}`jax.scipy.gaussian_kde` ({jax-issue}`#11237`).
jnp.ndarray: raise TypeError for binary operations with builtin collections
2022-06-24 10:03:27 -07:00
* Binary operations between JAX arrays and built-in collections (`dict`, `list`, `set`, `tuple`)
now raise a `TypeError` in all cases. Previously some cases (particularly equality and inequality)
would return boolean scalars inconsistent with similar operations in NumPy ({jax-issue}`#11234`).
Re-land #11498 after internal fixes. maintain an alias to `jax.tree_util.tree_map` in the top level `jax` module PiperOrigin-RevId: 463885774
2022-07-28 11:33:01 -07:00
* Several {mod}`jax.tree_util` routines accessed as top-level JAX package imports are now
deprecated, and will be removed in a future JAX release in accordance with the
{ref}`api-compatibility` policy:
Deprecate top-level aliases of jax.tree_util functions
2022-07-07 11:41:46 -07:00
* {func}`jax.treedef_is_leaf` is deprecated in favor of {func}`jax.tree_util.treedef_is_leaf`
* {func}`jax.tree_flatten` is deprecated in favor of {func}`jax.tree_util.tree_flatten`
* {func}`jax.tree_leaves` is deprecated in favor of {func}`jax.tree_util.tree_leaves`
* {func}`jax.tree_structure` is deprecated in favor of {func}`jax.tree_util.tree_structure`
* {func}`jax.tree_transpose` is deprecated in favor of {func}`jax.tree_util.tree_transpose`
* {func}`jax.tree_unflatten` is deprecated in favor of {func}`jax.tree_util.tree_unflatten`
jax.scipy.linalg.solve: deprecate the sym_pos argument following scipy 1.9.0
2022-07-19 13:57:49 -07:00
* The `sym_pos` argument of {func}`jax.scipy.linalg.solve` is deprecated in favor of `assume_a='pos'`,
following a similar deprecation in {func}`scipy.linalg.solve`.
Update versions for jax/jaxlib release
2022-06-21 12:41:26 -07:00
Update changelog for jax/jaxlib v0.3.15 release
2022-07-25 09:47:44 -07:00
## jaxlib 0.3.15 (July 22, 2022)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jaxlib-v0.3.14...jaxlib-v0.3.15).
Update versions for jax/jaxlib release
2022-06-21 12:41:26 -07:00
Update CHANGELOG.md and setup.py for 0.3.14 release.
2022-06-27 09:21:54 -04:00
## jax 0.3.14 (June 27, 2022)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.3.13...jax-v0.3.14).
Using etils(gfile) to support gcs buckets and file system for persistent compilation caching
2022-05-19 21:06:25 +00:00
* Breaking changes
Merge pull request #10771 from sshahrokhi:gfilecache PiperOrigin-RevId: 454692872
2022-06-13 13:58:15 -07:00
* {func}`jax.experimental.compilation_cache.initialize_cache` does not support
`max_cache_size_ bytes` anymore and will not get that as an input.
Change JAX_PLATFORMS to raise an exception when platform initialization fails
2022-06-17 16:38:56 +00:00
* `JAX_PLATFORMS` now raises an exception when platform initialization fails.
Add an implementation of jnp.linalg.slogdet based on QR decomposition. Adds a non-standard `method` argument to `jnp.linalg.slogdet` to select between the current LU decomposition based implementation (like NumPy) and the QR decomposition implementation. QR decomposition is more amenable to a high performance batched implementation particularly on TPU hardware because it does not need row pivoting. The same may be true on other hardware also, and having the option is nice either way! PiperOrigin-RevId: 449271317
2022-05-17 11:23:10 -07:00
* Changes
Update CHANGELOG.md and setup.py for 0.3.14 release.
2022-06-27 09:21:54 -04:00
* Fixed compatibility problems with NumPy 1.23.
Add an implementation of jnp.linalg.slogdet based on QR decomposition. Adds a non-standard `method` argument to `jnp.linalg.slogdet` to select between the current LU decomposition based implementation (like NumPy) and the QR decomposition implementation. QR decomposition is more amenable to a high performance batched implementation particularly on TPU hardware because it does not need row pivoting. The same may be true on other hardware also, and having the option is nice either way! PiperOrigin-RevId: 449271317
2022-05-17 11:23:10 -07:00
* {func}`jax.numpy.linalg.slogdet` now accepts an optional `method` argument
that allows selection between an LU-decomposition based implementation and
an implementation based on QR decomposition.
* {func}`jax.numpy.linalg.qr` now supports `mode="raw"`.
DeviceArray: Improve support for copy, deepcopy, and pickle
2022-05-17 13:20:38 -07:00
* `pickle`, `copy.copy`, and `copy.deepcopy` now have more complete support when
used on jax arrays ({jax-issue}`#10659`). In particular:
- `pickle` and `deepcopy` previously returned `np.ndarray` objects when used
on a `DeviceArray`; now `DeviceArray` objects are returned. For `deepcopy`,
the copied array is on the same device as the original. For `pickle` the
deserialized array will be on the default device.
- Within function transformations (i.e. traced code), `deepcopy` and `copy`
previously were no-ops. Now they use the same mechanism as `DeviceArray.copy()`.
Add jax.default_device to CHANGELOG
2022-06-03 17:20:09 -07:00
- Calling `pickle` on a traced array now results in an explicit
DeviceArray: Improve support for copy, deepcopy, and pickle
2022-05-17 13:20:38 -07:00
`ConcretizationTypeError`.
Update changelog to incorporate some recent changes.
2022-05-31 14:03:27 -04:00
* The implementation of singular value decomposition (SVD) and
symmetric/Hermitian eigendecomposition should be significantly faster on
TPU, especially for matrices above 1000x1000 or so. Both now use a spectral
divide-and-conquer algorithm for eigendecomposition (QDWH-eig).
[x64] jnp.ldexp: avoid implicit 64-bit promotion
2022-06-01 09:14:47 -07:00
* {func}`jax.numpy.ldexp` no longer silently promotes all inputs to float64,
instead it promotes to float32 for integer inputs of size int32 or smaller
({jax-issue}`#10921`).
Add an option to create a perfetto link in the JAX profiler
2022-05-25 16:01:16 -07:00
* Add a `create_perfetto_link` option to {func}`jax.profiler.start_trace` and
{func}`jax.profiler.start_trace`. When used, the profiler will generate a
link to the Perfetto UI to view the trace.
Store profiler server as a global variable and add a `stop_server` function
2022-05-31 15:42:46 -07:00
* Changed the semantics of {func}`jax.profiler.start_server(...)` to store the
keepalive globally, rather than requiring the user to keep a reference to
it.
Added random.generalized_normal and random.ball.
2022-06-03 15:11:29 -04:00
* Added {func}`jax.random.generalized_normal`.
* Added {func}`jax.random.ball`.
Add jax.default_device to CHANGELOG
2022-06-03 17:20:09 -07:00
* Added {func}`jax.default_device`.
Add collect_profile script
2022-06-02 22:15:53 -07:00
* Added a `python -m jax.collect_profile` script to manually capture program
Finalize deprecations of several APIs PiperOrigin-RevId: 633634215
2024-05-14 10:39:43 -07:00
traces as an alternative to the TensorBoard UI.
Update changelog for `named_scope` and adds it to the docs
2022-06-09 10:34:25 -07:00
* Added a `jax.named_scope` context manager that adds profiler metadata to
Python programs (similar to `jax.named_call`).
[x64] deprecate unsafe type casting in scatter-update operations
2022-06-09 15:21:49 -07:00
* In scatter-update operations (i.e. :attr:`jax.numpy.ndarray.at`), unsafe implicit
dtype casts are deprecated, and now result in a `FutureWarning`.
In a future release, this will become an error. An example of an unsafe implicit
cast is `jnp.zeros(4, dtype=int).at[0].set(1.5)`, in which `1.5` previously was
silently truncated to `1`.
Merge pull request #10771 from sshahrokhi:gfilecache PiperOrigin-RevId: 454692872
2022-06-13 13:58:15 -07:00
* {func}`jax.experimental.compilation_cache.initialize_cache` now supports gcs
Using etils(gfile) to support gcs buckets and file system for persistent compilation caching
2022-05-19 21:06:25 +00:00
bucket path as input.
Added scipy.stats.gennorm.
2022-06-11 14:24:19 -04:00
* Added {func}`jax.scipy.stats.gennorm`.
Update CHANGELOG.md and setup.py for 0.3.14 release.
2022-06-27 09:21:54 -04:00
* {func}`jax.numpy.roots` is now better behaved when `strip_zeros=False` when
coefficients have leading zeros ({jax-issue}`#11215`).
Update version after jax release PiperOrigin-RevId: 448822949
2022-05-15 12:13:54 -07:00
Update CHANGELOG.md and setup.py for 0.3.14 release.
2022-06-27 09:21:54 -04:00
## jaxlib 0.3.14 (June 27, 2022)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jaxlib-v0.3.10...jaxlib-v0.3.14).
Bump minimum Mac OS version to 10.14 (Mojave). It turns out that the support for C++17 is partial in 10.12, and in particular absl::optional and std::optional are not the same thing under 10.12. Increment to 10.14 which is the lowest version that builds successfully with absl::optional == std::optional. See: https://github.com/abseil/abseil-cpp/blob/89cdaed6557dbfb8cc8fba53a9dc8baf332df8b0/absl/base/config.h#L528 Strictly speaking, we could allow 10.13, but not without updating ABSL in the TF repository to incorporate https://github.com/abseil/abseil-cpp/commit/c86347d4cec43074e64e225a8753728f4bfc5ed6 which fixes the version detection test to permit 10.13 as well.
2022-06-01 20:27:42 -04:00
* x86-64 Mac wheels now require Mac OS 10.14 (Mojave) or newer. Mac OS 10.14
was released in 2018, so this should not be a very onerous requirement.
Update CHANGELOG.md and setup.py for 0.3.14 release.
2022-06-27 09:21:54 -04:00
* The bundled version of NCCL was updated to 2.12.12, fixing some deadlocks.
* The Python flatbuffers package is no longer a dependency of jaxlib.
Update version after jax release PiperOrigin-RevId: 448822949
2022-05-15 12:13:54 -07:00
Update version + CHANGELOG for jax 0.3.13 release
2022-05-16 12:17:07 -07:00
## jax 0.3.13 (May 16, 2022)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.3.12...jax-v0.3.13).
Update version + CHANGELOG for jax 0.3.13 release
2022-05-16 12:17:07 -07:00
Update values after jax release PiperOrigin-RevId: 448854487
2022-05-15 18:35:15 -07:00
## jax 0.3.12 (May 15, 2022)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.3.11...jax-v0.3.12).
Update values after jax release PiperOrigin-RevId: 448854487
2022-05-15 18:35:15 -07:00
* Changes
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* Fixes [#10717](https://github.com/jax-ml/jax/issues/10717).
Update values after jax release PiperOrigin-RevId: 448854487
2022-05-15 18:35:15 -07:00
Update version after jax release PiperOrigin-RevId: 448822949
2022-05-15 12:13:54 -07:00
## jax 0.3.11 (May 15, 2022)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.3.10...jax-v0.3.11).
Add a sort_eigenvalues option to lax.linalg.eigh(). An upcoming change to add a more scalable QDWH-based TPU symmetric eigendecomposition requires that we can obtain the TPU eigenvalues unsorted. The option already exists in XLA, so we simply need to plumb it through to the lax primitive. PiperOrigin-RevId: 448047584
2022-05-11 11:45:28 -07:00
* Changes
* {func}`jax.lax.eigh` now accepts an optional `sort_eigenvalues` argument
that allows users to opt out of eigenvalue sorting on TPU.
Change implementation of jax.scipy.linalg.polar() and jax._src.scipy.eigh to use the QDWH decomposition from jax._src.lax.qdwh. Remove jax._src.lax.polar. PiperOrigin-RevId: 448241206
2022-05-12 07:15:55 -07:00
* Deprecations
Change non-array arguments to jax.lax.linalg functions to be keyword-only arguments. PiperOrigin-RevId: 448066207
2022-05-11 13:06:19 -07:00
* Non-array arguments to functions in {mod}`jax.lax.linalg` are now marked
keyword-only. As a backward-compatibility step passing keyword-only
arguments positionally yields a warning, but in a future JAX release passing
keyword-only arguments positionally will fail.
However, most users should prefer to use {mod}`jax.numpy.linalg` instead.
Change implementation of jax.scipy.linalg.polar() and jax._src.scipy.eigh to use the QDWH decomposition from jax._src.lax.qdwh. Remove jax._src.lax.polar. PiperOrigin-RevId: 448241206
2022-05-12 07:15:55 -07:00
* {func}`jax.scipy.linalg.polar_unitary`, which was a JAX extension to the
scipy API, is deprecated. Use {func}`jax.scipy.linalg.polar` instead.
Finish jax release PiperOrigin-RevId: 441342919
2022-04-12 18:12:35 -07:00
jax and jaxlib release PiperOrigin-RevId: 446295827
2022-05-03 14:52:12 -07:00
## jax 0.3.10 (May 3, 2022)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.3.9...jax-v0.3.10).
jax and jaxlib release PiperOrigin-RevId: 446295827
2022-05-03 14:52:12 -07:00
## jaxlib 0.3.10 (May 3, 2022)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jaxlib-v0.3.7...jaxlib-v0.3.10).
jax and jaxlib release PiperOrigin-RevId: 446295827
2022-05-03 14:52:12 -07:00
* Changes
* [TF commit](https://github.com/tensorflow/tensorflow/commit/207d50d253e11c3a3430a700af478a1d524a779a)
fixes an issue in the MHLO canonicalizer that caused constant folding to
take a long time or crash for certain programs.
update version and changelog for pypi
2022-04-29 20:15:52 -07:00
Update the version numbers after JAX release. PiperOrigin-RevId: 446092433
2022-05-02 19:50:38 -07:00
## jax 0.3.9 (May 2, 2022)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.3.8...jax-v0.3.9).
Update the version numbers after JAX release. PiperOrigin-RevId: 446092433
2022-05-02 19:50:38 -07:00
* Changes
* Added support for fully asynchronous checkpointing for GlobalDeviceArray.
update version and changelog for pypi
2022-04-29 20:15:52 -07:00
## jax 0.3.8 (April 29 2022)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.3.7...jax-v0.3.8).
Delete the `mesh` context manager. The replacement for it is `Mesh`. PiperOrigin-RevId: 442619711
2022-04-18 13:37:16 -07:00
* Changes
[linalg] Add tpu svd lowering rule. PiperOrigin-RevId: 445533767
2022-04-29 16:42:08 -07:00
* {func}`jax.numpy.linalg.svd` on TPUs uses a qdwh-svd solver.
* {func}`jax.numpy.linalg.cond` on TPUs now accepts complex input.
* {func}`jax.numpy.linalg.pinv` on TPUs now accepts complex input.
* {func}`jax.numpy.linalg.matrix_rank` on TPUs now accepts complex input.
implement scipy.cluster.vq.vq also add no check_finite and overwrite_* docstring for some scipy.linalg functions
2022-04-20 22:38:50 +08:00
* {func}`jax.scipy.cluster.vq.vq` has been added.
Delete the `mesh` context manager. The replacement for it is `Mesh`. PiperOrigin-RevId: 442619711
2022-04-18 13:37:16 -07:00
* `jax.experimental.maps.mesh` has been deleted.
Please use `jax.experimental.maps.Mesh`. Please see https://jax.readthedocs.io/en/latest/_autosummary/jax.experimental.maps.Mesh.html#jax.experimental.maps.Mesh
for more information.
jax.scipy.qr: fix return type for mode='r'
2022-04-26 11:26:56 -07:00
* {func}`jax.scipy.linalg.qr` now returns a length-1 tuple rather than the raw array when
`mode='r'`, in order to match the behavior of `scipy.linalg.qr` ({jax-issue}`#10452`)
Add an optional mode= argument to jnp.take_along_axis. This allows users of jnp.take_along_axis to override the out-of-bounds indexing behavior. Default to "clip", which for the forward computation is identical to the current behavior. In a future change, we will change this to "fill".
2022-04-19 16:05:29 -04:00
* {func}`jax.numpy.take_along_axis` now takes an optional `mode` parameter
[JAX] Change jnp.take_along_axis to return invalid (e.g. NaN) values for out-of-bounds indices. Previously, out-of-bounds indices were clipped into range, but that behavior is error prone. We would rather fail in a more visible way when out-of-bounds indices are used. Future changes will migrate other JAX indexing operations to have the same semantics. PiperOrigin-RevId: 443390170
2022-04-21 08:51:35 -07:00
that specifies the behavior of out-of-bounds indexing. By default,
invalid values (e.g., NaN) will be returned for out-of-bounds indices. In
previous versions of JAX, invalid indices were clamped into range. The
previous behavior can be restored by passing `mode="clip"`.
Change the default jnp.take mode to "fill". Previously, `jnp.take` defaulted to clamping out-of-bounds indices into range. Now, `jnp.take` returns invalid values (e.g., NaN) for out-of-bounds indices. This change attempts to prevent latent bugs caused by inadvertent out-of-bounds indices. The previous behavior can be approximated using the "clip" or "wrap" fill modes. PiperOrigin-RevId: 445130143
2022-04-28 06:01:22 -07:00
* {func}`jax.numpy.take` now defaults to `mode="fill"`, which returns
invalid values (e.g., NaN) for out-of-bounds indices.
* Scatter operations, such as `x.at[...].set(...)`, now have `"drop"` semantics.
Change the default scatter mode to FILL_OR_DROP. This is a reasonably safe change, because it has no effect on the forward pass of a computation: the default behavior (PROMISE_IN_BOUNDS) also drops out-of-bounds scatters. This change does however affect the transpose (gradient) of a scatter with out-of-bounds indices: the gradient of a PROMISE_IN_BOUNDS scatter is a PROMISE_IN_BOUNDS gather, and a PROMISE_IN_BOUNDS gather clips out-of-bounds indices into range. This is not mathematically correct: a dropped scatter index does not contribute to the primal output, and so its transpose should yield a zero cotangent. After this change, the gradient of a default scatter is a gather with a fill value of 0: i.e., the indices that were dropped do not make gradient contributions, which is mathematically correct. Separately, I am working towards switching out-of-bounds gather() operations to also have FILL_OR_DROP semantics, although that change is more disruptive because a number of users have out-of-bounds indices in their gather()s. Issues: https://github.com/google/jax/issues/278 https://github.com/google/jax/issues/9839 PiperOrigin-RevId: 444935241
2022-04-27 12:26:17 -07:00
This has no effect on the scatter operation itself, but it means that when
differentiated the gradient of a scatter will yield zero cotangents for
out-of-bounds indices. Previously out-of-bounds indices were clamped into
range for the gradient, which was not mathematically correct.
Change jnp.take_along_axis to require that its indices are of integer type. Previously jnp.take_along_axis silently casted its indices to integers if they were not already integers. PiperOrigin-RevId: 443124521
2022-04-20 10:04:37 -07:00
* {func}`jax.numpy.take_along_axis` now raises a `TypeError` if its indices
are not of an integer type, matching the behavior of
{func}`numpy.take_along_axis`. Previously non-integer indices were silently
cast to integers.
replace int with operator.index part2 This change align the behavior of `ravel_multi_index`, `split` and `indices` to their `numpy` counterparts. Also ensure size argument of `nonzero` should be integer. The changes with `*space` are only simplification
2022-04-21 13:15:03 +08:00
* {func}`jax.numpy.ravel_multi_index` now raises a `TypeError` if its `dims` argument
is not of an integer type, matching the behavior of
{func}`numpy.ravel_multi_index`. Previously non-integer `dims` was silently
cast to integers.
* {func}`jax.numpy.split` now raises a `TypeError` if its `axis` argument
is not of an integer type, matching the behavior of
{func}`numpy.split`. Previously non-integer `axis` was silently
cast to integers.
* {func}`jax.numpy.indices` now raises a `TypeError` if its dimensions
are not of an integer type, matching the behavior of
{func}`numpy.indices`. Previously non-integer dimensions were silently
cast to integers.
* {func}`jax.numpy.diag` now raises a `TypeError` if its `k` argument
is not of an integer type, matching the behavior of
{func}`numpy.diag`. Previously non-integer `k` was silently
cast to integers.
Added random.orthogonal.
2022-04-29 14:20:50 -04:00
* Added {func}`jax.random.orthogonal`.
Deprecate remaining functionality in jax.test_util
2022-04-21 12:12:40 -07:00
* Deprecations
* Many functions and objects available in {mod}`jax.test_util` are now deprecated and will raise a
warning on import. This includes `cases_from_list`, `check_close`, `check_eq`, `device_under_test`,
`format_shape_dtype_string`, `rand_uniform`, `skip_on_devices`, `with_config`, `xla_bridge`, and
`_default_tolerance` ({jax-issue}`#10389`). These, along with previously-deprecated `JaxTestCase`,
`JaxTestLoader`, and `BufferDonationTestCase`, will be removed in a future JAX release.
Fix Typos in CHANGELOG.md This PR fixes the typos in Change log documentation
2024-03-12 13:57:07 +05:30
Most of these utilities can be replaced by calls to standard python & numpy testing utilities found
CHANGELOG: update test_util deprecation discussion
2022-04-21 13:37:56 -07:00
in e.g. {mod}`unittest`, {mod}`absl.testing`, {mod}`numpy.testing`, etc. JAX-specific functionality
such as device checking can be replaced through the use of public APIs such as {func}`jax.devices`.
Many of the deprecated utilities will still exist in {mod}`jax._src.test_util`, but these are not
public APIs and as such may be changed or removed without notice in future releases.
Deprecate DeviceArray.tile method
2022-04-13 12:11:10 -07:00
Jax 0.3.7 and jaxlib 0.3.7 release.
2022-04-15 12:02:05 -04:00
## jax 0.3.7 (April 15, 2022)
add random.loggamma and improve dirichlet & beta implementation
2022-03-21 08:33:11 -07:00
* [GitHub
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
commits](https://github.com/jax-ml/jax/compare/jax-v0.3.6...jax-v0.3.7).
Jax 0.3.7 and jaxlib 0.3.7 release.
2022-04-15 12:02:05 -04:00
* Changes:
* Fixed a performance problem if the indices passed to
{func}`jax.numpy.take_along_axis` were broadcasted ({jax-issue}`#10281`).
* {func}`jax.scipy.special.expit` and {func}`jax.scipy.special.logit` now
require their arguments to be scalars or JAX arrays. They also now promote
integer arguments to floating point.
2nd attempt at jax/jaxlib 0.3.7 release.
2022-04-15 15:20:02 -04:00
* The `DeviceArray.tile()` method is deprecated, because numpy arrays do not have a
`tile()` method. As a replacement for this, use {func}`jax.numpy.tile`
({jax-issue}`#10266`).
Update version numbers after jax/jaxlib release.
2022-04-07 16:38:09 -04:00
Fix typo in changelog.
2022-04-18 08:17:14 -04:00
## jaxlib 0.3.7 (April 15, 2022)
2nd attempt at jax/jaxlib 0.3.7 release.
2022-04-15 15:20:02 -04:00
* Changes:
* Linux wheels are now built conforming to the `manylinux2014` standard, instead
of `manylinux2010`.
Update version numbers after jax/jaxlib release.
2022-04-07 16:38:09 -04:00
Finish jax release PiperOrigin-RevId: 441342919
2022-04-12 18:12:35 -07:00
## jax 0.3.6 (April 12, 2022)
* [GitHub
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
commits](https://github.com/jax-ml/jax/compare/jax-v0.3.5...jax-v0.3.6).
Finish jax release PiperOrigin-RevId: 441342919
2022-04-12 18:12:35 -07:00
* Changes:
Jax 0.3.7 and jaxlib 0.3.7 release.
2022-04-15 12:02:05 -04:00
* Upgraded libtpu wheel to a version that fixes a hang when initializing a TPU
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
pod. Fixes [#10218](https://github.com/jax-ml/jax/issues/10218).
Deprecate jax.experimental.loops, step 2. Add deprecation warning and remove the tests. PiperOrigin-RevId: 441828243
2022-04-14 12:38:18 -07:00
* Deprecations:
Fix incorrect cross-reference breaking readthedocs build.
2022-04-14 16:12:48 -04:00
* {mod}`jax.experimental.loops` is being deprecated. See {jax-issue}`#10278`
Deprecate remaining functionality in jax.test_util
2022-04-21 12:12:40 -07:00
for an alternative API.
Update version numbers after jax/jaxlib release.
2022-04-07 16:38:09 -04:00
## jax 0.3.5 (April 7, 2022)
* [GitHub
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
commits](https://github.com/jax-ml/jax/compare/jax-v0.3.4...jax-v0.3.5).
add random.loggamma and improve dirichlet & beta implementation
2022-03-21 08:33:11 -07:00
* Changes:
* added {func}`jax.random.loggamma` & improved behavior of {func}`jax.random.beta`
Rename jax.nn.normalize to standardize. Add normalize alias with DeprecationWarning.
2022-03-23 20:39:39 +00:00
and {func}`jax.random.dirichlet` for small parameter values ({jax-issue}`#9906`).
Remove lax_numpy from jax.numpy namespace This is a private module that was inadvertently exported in the past.
2022-03-25 15:02:45 -07:00
* the private `lax_numpy` submodule is no longer exposed in the `jax.numpy` namespace ({jax-issue}`#10029`).
Implement jnp.from* array creation functions
2022-03-29 10:52:47 -07:00
* added array creation routines {func}`jax.numpy.frombuffer`, {func}`jax.numpy.fromfunction`,
and {func}`jax.numpy.fromstring` ({jax-issue}`#10049`).
Add CHANGELOG entry for #10069
2022-03-30 08:05:34 -07:00
* `DeviceArray.copy()` now returns a `DeviceArray` rather than a `np.ndarray` ({jax-issue}`#10069`)
Add jax.scipy.linalg.rsf2csf
2022-03-27 12:31:12 +01:00
* added {func}`jax.scipy.linalg.rsf2csf`
Update the changelog to say sharded_jit is deprecated.
2022-04-18 13:48:19 -07:00
* `jax.experimental.sharded_jit` has been deprecated and will be removed soon.
Rename jax.nn.normalize to standardize. Add normalize alias with DeprecationWarning.
2022-03-23 20:39:39 +00:00
* Deprecations:
* {func}`jax.nn.normalize` is being deprecated. Use {func}`jax.nn.standardize` instead ({jax-issue}`#9899`).
Deprecate jax.tree_util.tree_multimap
2022-04-01 14:51:54 -07:00
* {func}`jax.tree_util.tree_multimap` is deprecated. Use {func}`jax.tree_util.tree_map` instead ({jax-issue}`#5746`).
Jax 0.3.5 and jaxlib 0.3.5 release.
2022-04-06 23:56:41 +00:00
* `jax.experimental.sharded_jit` is deprecated. Use `pjit` instead.
update version and changelog for pypi
2022-03-17 15:35:26 -07:00
Update version numbers after jax/jaxlib release.
2022-04-07 16:38:09 -04:00
## jaxlib 0.3.5 (April 7, 2022)
[GPU] Force an input buffer copy for double precision complex-to-real IRFFTs. Fixes https://github.com/google/jax/issues/9946 PiperOrigin-RevId: 439414091
2022-04-04 14:34:00 -07:00
* Bug fixes
* Fixed a bug where double-precision complex-to-real IRFFTs would mutate their
input buffers on GPU ({jax-issue}`#9946`).
Jax 0.3.5 and jaxlib 0.3.5 release.
2022-04-06 23:56:41 +00:00
* Fixed incorrect constant-folding of complex scatters ({jax-issue}`#10159`)
update version and changelog for pypi
2022-03-17 15:35:26 -07:00
update version and changelog for pypi
2022-03-18 14:16:00 -07:00
## jax 0.3.4 (March 18, 2022)
* [GitHub
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
commits](https://github.com/jax-ml/jax/compare/jax-v0.3.3...jax-v0.3.4).
update version and changelog for pypi
2022-03-18 14:16:00 -07:00
update version and changelog for pypi
2022-03-17 15:35:26 -07:00
## jax 0.3.3 (March 17, 2022)
Update values after release PiperOrigin-RevId: 427910510
2022-02-10 20:32:20 -08:00
* [GitHub
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
commits](https://github.com/jax-ml/jax/compare/jax-v0.3.2...jax-v0.3.3).
Update CHANGELOG and setup.py for jax + jaxlib 0.3.2 releases
2022-03-16 10:17:42 -07:00
## jax 0.3.2 (March 16, 2022)
* [GitHub
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
commits](https://github.com/jax-ml/jax/compare/jax-v0.3.1...jax-v0.3.2).
Remove jax.ops.index... functions. These functions have been deprecated and have issued a DeprecationWarning since jax 0.2.22 in October 2021.
2022-02-17 15:23:44 -05:00
* Changes:
* The functions `jax.ops.index_update`, `jax.ops.index_add`, which were
deprecated in 0.2.22, have been removed. Please use
[the `.at` property on JAX arrays](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html)
instead, e.g., `x.at[idx].set(y)`.
[JAX] Add release note for the graduration of the experimental.ann module. PiperOrigin-RevId: 431951602
2022-03-02 08:58:22 -08:00
* Moved `jax.experimental.ann.approx_*_k` into `jax.lax`. These functions are
optimized alternatives to `jax.lax.top_k`.
Call _check_arraylike on inputs to broadcast_to and broadcast_arrays
2022-03-02 15:23:34 -08:00
* {func}`jax.numpy.broadcast_arrays` and {func}`jax.numpy.broadcast_to` now require scalar
or array-like inputs, and will fail if they are passed lists (part of {jax-issue}`#7737`).
Update CHANGELOG and setup.py for jax + jaxlib 0.3.2 releases
2022-03-16 10:17:42 -07:00
* The standard jax[tpu] install can now be used with Cloud TPU v4 VMs.
Bump jax and jaxlib versions for 0.3.2 release Also add CPU pjit to changelog
2022-03-16 14:25:19 -07:00
* `pjit` now works on CPU (in addition to previous TPU and GPU support).
Remove jax.ops.index... functions. These functions have been deprecated and have issued a DeprecationWarning since jax 0.2.22 in October 2021.
2022-02-17 15:23:44 -05:00
Finish jax release PiperOrigin-RevId: 429670894
2022-02-18 16:23:02 -08:00
Update CHANGELOG and setup.py for jax + jaxlib 0.3.2 releases
2022-03-16 10:17:42 -07:00
## jaxlib 0.3.2 (March 16, 2022)
Finish jax release PiperOrigin-RevId: 429670894
2022-02-18 16:23:02 -08:00
* Changes
Add boolean flag to as_hlo_text to enable writing large constants. PiperOrigin-RevId: 434556535
2022-03-14 13:45:49 -07:00
* ``XlaComputation.as_hlo_text()`` now supports printing large constants by
passing boolean flag ``print_large_constants=True``.
[JAX] Deprecate .block_host_until_ready() in favor of .block_until_ready(). JAX kept an older name around (.block_host_until_ready()) in parallel with the new name (.block_until_ready()) to avoid breaking users. Deprecate it so we only have one name. PiperOrigin-RevId: 433228545
2022-03-08 09:14:14 -08:00
* Deprecations:
* The ``.block_host_until_ready()`` method on JAX arrays has been deprecated.
Use ``.block_until_ready()`` instead.
Finish jax release PiperOrigin-RevId: 429670894
2022-02-18 16:23:02 -08:00
## jax 0.3.1 (Feb 18, 2022)
* [GitHub
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
commits](https://github.com/jax-ml/jax/compare/jax-v0.3.0...jax-v0.3.1).
Add deprecation warning to JaxTestCase and JaxTestLoader
2022-02-17 14:58:58 -08:00
* Changes:
* `jax.test_util.JaxTestCase` and `jax.test_util.JaxTestLoader` are now deprecated.
The suggested replacement is to use `parametrized.TestCase` directly. For tests that
rely on custom asserts such as `JaxTestCase.assertAllClose()`, the suggested replacement
is to use standard numpy testing utilities such as {func}`numpy.testing.assert_allclose()`,
which work directly with JAX arrays ({jax-issue}`#9620`).
* `jax.test_util.JaxTestCase` now sets `jax_numpy_rank_promotion='raise'` by default
({jax-issue}`#9562`). To recover the previous behavior, use the new
`jax.test_util.with_config` decorator:
```python
@jtu.with_config(jax_numpy_rank_promotion='allow')
class MyTestCase(jtu.JaxTestCase):
...
```
* Added {func}`jax.scipy.linalg.schur`, {func}`jax.scipy.linalg.sqrtm`,
{func}`jax.scipy.signal.csd`, {func}`jax.scipy.signal.stft`,
{func}`jax.scipy.signal.welch`.
Update values after release PiperOrigin-RevId: 427910510
2022-02-10 20:32:20 -08:00
Release jax and jaxlib 0.3.0 as per the new release process. PiperOrigin-RevId: 427809845
2022-02-10 11:58:30 -08:00
## jax 0.3.0 (Feb 10, 2022)
Bump jax version to 0.2.29 and update CHANGELOG
2022-02-01 17:59:49 -08:00
* [GitHub
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
commits](https://github.com/jax-ml/jax/compare/jax-v0.2.28...jax-v0.3.0).
Bump jax version to 0.2.29 and update CHANGELOG
2022-02-01 17:59:49 -08:00
Release jax and jaxlib 0.3.0 as per the new release process. PiperOrigin-RevId: 427809845
2022-02-10 11:58:30 -08:00
* Changes
* jax version has been bumped to 0.3.0. Please see the [design doc](https://jax.readthedocs.io/en/latest/design_notes/jax_versioning.html)
for the explanation.
## jaxlib 0.3.0 (Feb 10, 2022)
* Changes
* Bazel 5.0.0 is now required to build jaxlib.
* jaxlib version has been bumped to 0.3.0. Please see the [design doc](https://jax.readthedocs.io/en/latest/design_notes/jax_versioning.html)
for the explanation.
## jax 0.2.28 (Feb 1, 2022)
Bump jax version after jax release PiperOrigin-RevId: 415064518
2021-12-08 12:07:42 -08:00
* [GitHub
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
commits](https://github.com/jax-ml/jax/compare/jax-v0.2.27...jax-v0.2.28).
Change the default IR dialect returned by .compiler_ir() to MHLO. PiperOrigin-RevId: 423091674
2022-01-20 09:49:40 -08:00
* `jax.jit(f).lower(...).compiler_ir()` now defaults to the MHLO dialect if no
`dialect=` is passed.
[JAX] Change signature of .mhlo() method on compiler IR objects to return an ir.Module object instead of its string representation. It isn't free to pretty-print IR, so it's best to avoid it unless necessary. In addition, by returning an IR object, the user is now free to, say, print it with different options. For example, one can now write things like: ``` In [1]: import numpy as np, jax, jax.numpy as jnp In [2]: m = jax.jit(lambda x: x + jnp.array(np.arange(1000))).lower(7.).compiler_ir(dialect='mhlo') In [3]: m.operation.print(large_elements_limit=10) module @jit__lambda_.4 { func public @main(%arg0: tensor<f32>) -> tensor<1000xf32> { %0 = mhlo.constant opaque<"_", "0xDEADBEEF"> : tensor<1000xi32> %1 = "mhlo.convert"(%0) : (tensor<1000xi32>) -> tensor<1000xf32> %2 = "mhlo.broadcast_in_dim"(%arg0) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1000xf32> %3 = mhlo.add %2, %1 : tensor<1000xf32> return %3 : tensor<1000xf32> } } ``` Fixes https://github.com/google/jax/issues/9226 PiperOrigin-RevId: 422855649
2022-01-19 11:01:03 -08:00
* The `jax.jit(f).lower(...).compiler_ir(dialect='mhlo')` now returns an MLIR
`ir.Module` object instead of its string representation.
update version and changelog for pypi
2022-01-18 11:38:32 -08:00
Update XLA for jaxlib 0.1.76 release.
2022-01-26 22:19:11 +00:00
## jaxlib 0.1.76 (Jan 27, 2022)
update version and changelog for pypi
2022-01-18 11:38:32 -08:00
* New features
* Includes precompiled SASS for NVidia compute capability 8.0 GPUS
(e.g. A100). Removes precompiled SASS for compute capability 6.1 so as not
to increase the number of compute capabilities: GPUs with compute capability
6.1 can use the 6.0 SASS.
Update XLA for jaxlib 0.1.76 release.
2022-01-26 22:19:11 +00:00
* With jaxlib 0.1.76, JAX uses the MHLO MLIR dialect as its primary target compiler IR
by default.
update version and changelog for pypi
2022-01-18 11:38:32 -08:00
* Breaking changes
* Support for NumPy 1.18 has been dropped, per the
[deprecation policy](https://jax.readthedocs.io/en/latest/deprecation.html).
Please upgrade to a supported NumPy version.
Update XLA for jaxlib 0.1.76 release.
2022-01-26 22:19:11 +00:00
* Bug fixes
* Fixed a bug where apparently identical pytreedef objects constructed by different routes
do not compare as equal (#9066).
* The JAX jit cache requires two static arguments to have identical types for a cache hit (#9311).
update version and changelog for pypi
2022-01-18 11:38:32 -08:00
## jax 0.2.27 (Jan 18 2022)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.2.26...jax-v0.2.27).
Bump jax version after jax release PiperOrigin-RevId: 415064518
2021-12-08 12:07:42 -08:00
[hcb] Simplifications to the host_calback API * dropping support for special AD handling for hcb.id_tap and id_print. From now on, only the primals are tapped. The old behavior can be obtained (for a limited time) by setting the JAX_HOST_CALLBACK_AD_TRANSFORMS environment variale, or the --flax_host_callback_ad_transforms flag. Additionally, added documentation for how to implement the old behavior using JAX custom AD APIs. This allows us to make some significant cleanup in the internals.
2021-11-24 12:58:16 +02:00
* Breaking changes:
Drop support for NumPy 1.18. Per NEP-29, we can drop NumPy 1.18 support on Dec 22, 2021. The next NumPy deprecation will be 1.19 on Jun 21, 2022. PiperOrigin-RevId: 419651428
2022-01-04 12:11:00 -08:00
* Support for NumPy 1.18 has been dropped, per the
[deprecation policy](https://jax.readthedocs.io/en/latest/deprecation.html).
Please upgrade to a supported NumPy version.
[hcb] Simplifications to the host_calback API * dropping support for special AD handling for hcb.id_tap and id_print. From now on, only the primals are tapped. The old behavior can be obtained (for a limited time) by setting the JAX_HOST_CALLBACK_AD_TRANSFORMS environment variale, or the --flax_host_callback_ad_transforms flag. Additionally, added documentation for how to implement the old behavior using JAX custom AD APIs. This allows us to make some significant cleanup in the internals.
2021-11-24 12:58:16 +02:00
* The host_callback primitives have been simplified to drop the
[hcb] Add support for remat2 to host_callback A callback under ad_checkpoint.checkpoint will be invoked twice when taking the gradient: once during the forward pass and once again during the backward pass when the residuals for the forward pass are rematerialized.
2021-11-24 12:58:16 +02:00
special autodiff handling for hcb.id_tap and id_print.
From now on, only the primals are tapped. The old behavior can be
obtained (for a limited time) by setting the ``JAX_HOST_CALLBACK_AD_TRANSFORMS``
[callback] Add a flag to implement host_callback in terms of io_callback. The host_callbacks APIs are deprecated and will be removed. In order to help the transition to the new APIs, we add a flag (`JAX_HOST_CALLBACK_LEGACY`) that when set to `False` will use `io_callback` (and `pure_callback` and `jax.debug.callback`) to implement the host_callback APIs. See issue #20385 for more details. We change the tests to accomodate slightly different results when using the new callbacks. The tests that use `tap_with_device` and `call_with_device` are disabled when using the new callbacks.
2024-03-29 13:36:20 +02:00
environment variable, or the ```--jax_host_callback_ad_transforms``` flag.
[hcb] Add support for remat2 to host_callback A callback under ad_checkpoint.checkpoint will be invoked twice when taking the gradient: once during the forward pass and once again during the backward pass when the residuals for the forward pass are rematerialized.
2021-11-24 12:58:16 +02:00
Additionally, added documentation for how to implement the old behavior
using JAX custom AD APIs ({jax-issue}`#8678`).
lax.sort: regularize handling of -0.0 and -NaN
2022-01-13 13:03:41 -08:00
* Sorting now matches the behavior of NumPy for ``0.0`` and ``NaN`` regardless of the
bit representation. In particular, ``0.0`` and ``-0.0`` are now treated as equivalent,
where previously ``-0.0`` was treated as less than ``0.0``. Additionally all ``NaN``
representations are now treated as equivalent and sorted to the end of the array.
Previously negative ``NaN`` values were sorted to the front of the array, and ``NaN``
values with different internal bit representations were not treated as equivalent, and
were sorted according to those bit patterns ({jax-issue}`#9178`).
jnp.unique: properly handle NaN values
2022-01-13 15:54:07 -08:00
* {func}`jax.numpy.unique` now treats ``NaN`` values in the same way as `np.unique` in
NumPy versions 1.21 and newer: at most one ``NaN`` value will appear in the uniquified
output ({jax-issue}`9184`).
[hcb] Add support for remat2 to host_callback A callback under ad_checkpoint.checkpoint will be invoked twice when taking the gradient: once during the forward pass and once again during the backward pass when the residuals for the forward pass are rematerialized.
2021-11-24 12:58:16 +02:00
* Bug fixes:
* host_callback now supports ad_checkpoint.checkpoint ({jax-issue}`#8907`).
[hcb] Simplifications to the host_calback API * dropping support for special AD handling for hcb.id_tap and id_print. From now on, only the primals are tapped. The old behavior can be obtained (for a limited time) by setting the JAX_HOST_CALLBACK_AD_TRANSFORMS environment variale, or the --flax_host_callback_ad_transforms flag. Additionally, added documentation for how to implement the old behavior using JAX custom AD APIs. This allows us to make some significant cleanup in the internals.
2021-11-24 12:58:16 +02:00
add jax.block_until_ready to docs and changelog also unrelatedly fix a couple of the uses of rst in changelog.md (though many others remain)
2021-12-14 13:36:46 -08:00
* New features:
* add `jax.block_until_ready` ({jax-issue}`#8941)
Add debugging flag for dumping the JAX-generated MHLO/HLO IR to a file. While HLO dumping is redundant with XLA's XLA_FLAGS=--xla_dump_to=... feature, MHLO dumping is useful since XLA only ever sees and dumps the IR after it has been canonicalized and converted to HLO. Some debugging tasks require easy access to the MHLO as well. PiperOrigin-RevId: 416435598
2021-12-14 17:43:40 -08:00
* Added a new debugging flag/environment variable `JAX_DUMP_IR_TO=/path`.
If set, JAX dumps the MHLO/HLO IR it generates for each computation to a
file under the given path.
Copybara import of the project: -- 4fcdadbfb3f4c484fd4432203cf13b88782b9311 by Matthew Johnson <mattjj@google.com>: add jax.ensure_compile_time_eval to public api aka jax.core.eval_context COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/7987 from google:issue7535 4fcdadbfb3f4c484fd4432203cf13b88782b9311 PiperOrigin-RevId: 420928687
2022-01-10 20:57:56 -08:00
* Added `jax.ensure_compile_time_eval` to the public api ({jax-issue}`#7987`).
[jax2tf] Add jax2tf_associative_scan_reductions flag This flag allows users to match the JAX performance for associative reductions in CPU. See README.md for details.
2022-01-13 12:28:25 +02:00
* jax2tf now supports a flag jax2tf_associative_scan_reductions to change
the lowering for associative reductions, e.g., jnp.cumsum, to behave
like JAX on CPU and GPU (to use an associative scan). See the jax2tf README
for more details ({jax-issue}`#9189`).
add jax.block_until_ready to docs and changelog also unrelatedly fix a couple of the uses of rst in changelog.md (though many others remain)
2021-12-14 13:36:46 -08:00
Update jaxlib version number to 0.1.76 PiperOrigin-RevId: 415050863
2021-12-08 11:13:37 -08:00
Bump jax version after jax release PiperOrigin-RevId: 415064518
2021-12-08 12:07:42 -08:00
## jaxlib 0.1.75 (Dec 8, 2021)
Update jaxlib version number to 0.1.76 PiperOrigin-RevId: 415050863
2021-12-08 11:13:37 -08:00
* New features:
* Support for python 3.10.
Update jaxlib versions.
2021-11-17 11:43:19 -05:00
Bump jax version after jax release PiperOrigin-RevId: 415064518
2021-12-08 12:07:42 -08:00
## jax 0.2.26 (Dec 8, 2021)
update version and changelog for pypi
2021-11-10 13:41:25 -08:00
* [GitHub
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
commits](https://github.com/jax-ml/jax/compare/jax-v0.2.25...jax-v0.2.26).
[jax2tf] Force TF compilation for code under jax.jit. Previously, jax.jit was ignored by jax2tf. This can result in the converted code being much slower than the JAX core, unless the user adds an explicit `tf.function(jit_compile=True)`. With this change that wrapper is added automatically for all code fragments under jax.jit. Note that most jax.numpy functions are annotated with jax.jit, so with this change they will all be compiled. When doing this I ran into problems with tf.custom_gradient and tf.function. As documented in the [tf.custom_gradient](https://www.tensorflow.org/api_docs/python/tf/custom_gradient) documentation, you get a LookupError when trying to build the gradient of a tf.function, even if it has a tf.custom_gradient defined. The recommended solution is to add a tf.stop_gradient. This is safe, since jax2tf will always wrap the converted functions with a tf.custom_gradient.
2021-09-07 18:36:47 +03:00
Change the default out-of-bounds behavior for jax.ops.segment_... to FILL_OR_DROP. This matches the documented behavior. Fixes https://github.com/google/jax/issues/8634 PiperOrigin-RevId: 411635687
2021-11-22 13:32:25 -08:00
* Bug fixes:
* Out-of-bounds indices to `jax.ops.segment_sum` will now be handled with
Fix Typos in CHANGELOG.md This PR fixes the typos in Change log documentation
2024-03-12 13:57:07 +05:30
`FILL_OR_DROP` semantics, as documented. This primarily affects the
Change the default out-of-bounds behavior for jax.ops.segment_... to FILL_OR_DROP. This matches the documented behavior. Fixes https://github.com/google/jax/issues/8634 PiperOrigin-RevId: 411635687
2021-11-22 13:32:25 -08:00
reverse-mode derivative, where gradients corresponding to out-of-bounds
indices will now be returned as 0. (#8634).
[jax2tf] Force TF compilation for code under jax.jit. Previously, jax.jit was ignored by jax2tf. This can result in the converted code being much slower than the JAX core, unless the user adds an explicit `tf.function(jit_compile=True)`. With this change that wrapper is added automatically for all code fragments under jax.jit. Note that most jax.numpy functions are annotated with jax.jit, so with this change they will all be compiled. When doing this I ran into problems with tf.custom_gradient and tf.function. As documented in the [tf.custom_gradient](https://www.tensorflow.org/api_docs/python/tf/custom_gradient) documentation, you get a LookupError when trying to build the gradient of a tf.function, even if it has a tf.custom_gradient defined. The recommended solution is to add a tf.stop_gradient. This is safe, since jax2tf will always wrap the converted functions with a tf.custom_gradient.
2021-09-07 18:36:47 +03:00
* jax2tf will force the converted code to use XLA for the code fragments
under jax.jit, e.g., most jax.numpy functions ({jax-issue}`#7839`).
update version and changelog for pypi
2021-11-10 13:41:25 -08:00
Update jaxlib versions.
2021-11-17 11:43:19 -05:00
## jaxlib 0.1.74 (Nov 17, 2021)
Fix test failures and update changelog. Use dtypes.issubdtype to test for subtyping otherwise we mishandle bfloat16 dtypes. Don't pass an empty list to concatenate() when converting a shape to a value. Forbid empty lists as arguments to lax.concatenate().
2021-11-16 17:36:28 -05:00
* Enabled peer-to-peer copies between GPUs. Previously, GPU copies were bounced via
the host, which is usually slower.
* Added experimental MLIR Python bindings for use by JAX.
Add a default cuda installation path and more explicit installation paths for CUDA jaxlib. ``` # Installs Cuda 11 with Cudnn 8.2 $ pip install jax[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html $ pip install jax[cuda=11,cudnn=82] -f https://storage.googleapis.com/jax-releases/jax_releases.html $ pip install jax[cuda=11,cudnn=805] -f https://storage.googleapis.com/jax-releases/jax_releases.html ``` PiperOrigin-RevId: 404134291
2021-10-18 19:55:51 -07:00
update version and changelog for pypi
2021-11-10 13:41:25 -08:00
## jax 0.2.25 (Nov 10, 2021)
Update jax version and CHANGELOG for 0.2.22 release
2021-10-12 18:21:11 -07:00
* [GitHub
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
commits](https://github.com/jax-ml/jax/compare/jax-v0.2.24...jax-v0.2.25).
Update jax version number for jax release. PiperOrigin-RevId: 404262742
2021-10-19 08:04:59 -07:00
Add jax.distributed.initialize for multi-host GPU.
2021-10-25 15:34:57 -07:00
* New features:
* (Experimental) `jax.distributed.initialize` exposes multi-host GPU backend.
jax.random.permutation: add independent keyword
2021-11-02 11:39:41 -07:00
* `jax.random.permutation` supports new `independent` keyword argument
({jax-issue}`#8430`)
[JAX] move example libraries from `jax.experimental` into `jax.example_libraries` The `jax.experimental.stax` and `jax.experimental.optimizers` modules are standalone examples libraries. By contrast, the remaining modules in `jax.experimental` are experimental features of the JAX core system. This change moves the two example libraries, and the README that describes them, to `jax.example_libraries` to reflect this distinction. PiperOrigin-RevId: 404405186
2021-10-19 17:30:16 -07:00
* Breaking changes
* Moved `jax.experimental.stax` to `jax.example_libraries.stax`
* Moved `jax.experimental.optimizers` to `jax.example_libraries.optimizers`
[JAX] Added `jax.lax.linalg.qdwh`. PiperOrigin-RevId: 406453671
2021-10-29 14:44:27 -07:00
* New features:
* Added `jax.lax.linalg.qdwh`.
[JAX] move example libraries from `jax.experimental` into `jax.example_libraries` The `jax.experimental.stax` and `jax.experimental.optimizers` modules are standalone examples libraries. By contrast, the remaining modules in `jax.experimental` are experimental features of the JAX core system. This change moves the two example libraries, and the README that describes them, to `jax.example_libraries` to reflect this distinction. PiperOrigin-RevId: 404405186
2021-10-19 17:30:16 -07:00
Update jax version number for jax release. PiperOrigin-RevId: 404262742
2021-10-19 08:04:59 -07:00
## jax 0.2.24 (Oct 19, 2021)
* [GitHub
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
commits](https://github.com/jax-ml/jax/compare/jax-v0.2.22...jax-v0.2.24).
Update jax version and CHANGELOG for 0.2.22 release
2021-10-12 18:21:11 -07:00
Fix CHANGELOG.md
2021-10-13 17:11:50 -06:00
* New features:
* `jax.random.choice` and `jax.random.permutation` now support
multidimensional arrays and an optional `axis` argument ({jax-issue}`#8158`)
jnp.take/jnp.take_along_axis: require array inputs
2021-10-14 11:54:28 -07:00
* Breaking changes:
* `jax.numpy.take` and `jax.numpy.take_along_axis` now require array-like inputs
(see {jax-issue}`#7737`)
Fix CHANGELOG.md
2021-10-13 17:11:50 -06:00
Add a default cuda installation path and more explicit installation paths for CUDA jaxlib. ``` # Installs Cuda 11 with Cudnn 8.2 $ pip install jax[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html $ pip install jax[cuda=11,cudnn=82] -f https://storage.googleapis.com/jax-releases/jax_releases.html $ pip install jax[cuda=11,cudnn=805] -f https://storage.googleapis.com/jax-releases/jax_releases.html ``` PiperOrigin-RevId: 404134291
2021-10-18 19:55:51 -07:00
## jaxlib 0.1.73 (Oct 18, 2021)
* Multiple cuDNN versions are now supported for jaxlib GPU `cuda11` wheels.
* cuDNN 8.2 or newer. We recommend using the cuDNN 8.2 wheel if your cuDNN
installation is new enough, since it supports additional functionality.
* cuDNN 8.0.5 or newer.
* Breaking changes:
* The install commands for GPU jaxlib are as follows:
```bash
pip install --upgrade pip
# Installs the wheel compatible with CUDA 11 and cuDNN 8.2 or newer.
Use `cuda11_cudnn82` instead of `cuda=11,cudnn=82` because the latter one is a syntax error PiperOrigin-RevId: 404240654
2021-10-19 06:24:27 -07:00
pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html
Add a default cuda installation path and more explicit installation paths for CUDA jaxlib. ``` # Installs Cuda 11 with Cudnn 8.2 $ pip install jax[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html $ pip install jax[cuda=11,cudnn=82] -f https://storage.googleapis.com/jax-releases/jax_releases.html $ pip install jax[cuda=11,cudnn=805] -f https://storage.googleapis.com/jax-releases/jax_releases.html ``` PiperOrigin-RevId: 404134291
2021-10-18 19:55:51 -07:00
# Installs the wheel compatible with Cuda 11 and cudnn 8.2 or newer.
Use `cuda11_cudnn82` instead of `cuda=11,cudnn=82` because the latter one is a syntax error PiperOrigin-RevId: 404240654
2021-10-19 06:24:27 -07:00
pip install jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_releases.html
Add a default cuda installation path and more explicit installation paths for CUDA jaxlib. ``` # Installs Cuda 11 with Cudnn 8.2 $ pip install jax[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html $ pip install jax[cuda=11,cudnn=82] -f https://storage.googleapis.com/jax-releases/jax_releases.html $ pip install jax[cuda=11,cudnn=805] -f https://storage.googleapis.com/jax-releases/jax_releases.html ``` PiperOrigin-RevId: 404134291
2021-10-18 19:55:51 -07:00
# Installs the wheel compatible with Cuda 11 and cudnn 8.0.5 or newer.
Use `cuda11_cudnn82` instead of `cuda=11,cudnn=82` because the latter one is a syntax error PiperOrigin-RevId: 404240654
2021-10-19 06:24:27 -07:00
pip install jax[cuda11_cudnn805] -f https://storage.googleapis.com/jax-releases/jax_releases.html
Add a default cuda installation path and more explicit installation paths for CUDA jaxlib. ``` # Installs Cuda 11 with Cudnn 8.2 $ pip install jax[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html $ pip install jax[cuda=11,cudnn=82] -f https://storage.googleapis.com/jax-releases/jax_releases.html $ pip install jax[cuda=11,cudnn=805] -f https://storage.googleapis.com/jax-releases/jax_releases.html ``` PiperOrigin-RevId: 404134291
2021-10-18 19:55:51 -07:00
```
Remove 10.2 cuda support PiperOrigin-RevId: 402707900
2021-10-12 18:43:40 -07:00
Update jax version and CHANGELOG for 0.2.22 release
2021-10-12 18:21:11 -07:00
## jax 0.2.22 (Oct 12, 2021)
Add Changelog for jax v0.2.21 development
2021-09-15 12:10:30 -07:00
* [GitHub
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
commits](https://github.com/jax-ml/jax/compare/jax-v0.2.21...jax-v0.2.22).
Disallow non-hashable static arguments in pmap(). * Don't wrap static arguments in hashable wrappers in pmap. * Delete wrap_hashably(). * In argnums_partial, either enforce hashability or wrap values with an explicitly unhashable wrapper. The intent here is that either we should check for hashability early or we should make sure it's clear that it's not something we intended.. * Delete argnames_partial, which appears unused.
2021-07-19 13:11:38 -04:00
* Breaking Changes
* Static arguments to `jax.pmap` must now be hashable.
Unhashable static arguments have long been disallowed on `jax.jit`, but they
were still permitted on `jax.pmap`; `jax.pmap` compared unhashable static
arguments using object identity.
This behavior is a footgun, since comparing arguments using
object identity leads to recompilation each time the object identity
changes. Instead, we now ban unhashable arguments: if a user of `jax.pmap`
wants to compare static arguments by object identity, they can define
`__hash__` and `__eq__` methods on their objects that do that, or wrap their
objects in an object that has those operations with object identity
semantics. Another option is to use `functools.partial` to encapsulate the
unhashable static arguments into the function object.
* `jax.util.partial` was an accidental export that has now been removed. Use
`functools.partial` from the Python standard library instead.
Add note to changelog about deprecation of jax.ops.index_...
2021-10-06 09:31:42 -04:00
* Deprecations
* The functions `jax.ops.index_update`, `jax.ops.index_add` etc. are
deprecated and will be removed in a future JAX release. Please use
[the `.at` property on JAX arrays](https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html)
instead, e.g., `x.at[idx].set(y)`. For now, these functions produce a
`DeprecationWarning`.
Enable C++ pmap. On CPU: ``` name old cpu/op new cpu/op delta pmap_trivial_2_devices 128µs ± 6% 14µs ± 3% -89.06% (p=0.008 n=5+5) pmap_trivial_dispatch_8_devices 212µs ± 2% 35µs ± 1% -83.54% (p=0.008 n=5+5) pmap_trivial_8_devices 215µs ± 1% 40µs ± 4% -81.31% (p=0.008 n=5+5) pmap_simple_2_devices 123µs ± 5% 15µs ± 6% -87.70% (p=0.008 n=5+5) pmap_simple_dispatch_8_devices 211µs ± 3% 35µs ± 2% -83.24% (p=0.008 n=5+5) pmap_simple_8_devices 217µs ± 5% 40µs ± 2% -81.68% (p=0.008 n=5+5) pmap_simple_dispatch_8_devices_100_args 5.42ms ± 7% 0.52ms ± 2% -90.44% (p=0.008 n=5+5) pmap_simple_8_devices_100_args 26.5ms ±21% 17.5ms ±37% -34.18% (p=0.008 n=5+5) sda_index_1 7.45µs ± 6% 7.53µs ± 6% ~ (p=0.222 n=5+5) sda_index_2 14.1µs ± 1% 14.3µs ± 4% ~ (p=0.690 n=5+5) sda_index_8 56.0µs ± 3% 56.9µs ± 4% ~ (p=0.310 n=5+5) name old time/op new time/op delta pmap_trivial_2_devices 136µs ± 8% 19µs ± 3% -86.08% (p=0.008 n=5+5) pmap_trivial_dispatch_8_devices 216µs ± 3% 39µs ± 2% -81.94% (p=0.008 n=5+5) pmap_trivial_8_devices 219µs ± 2% 49µs ±38% -77.67% (p=0.008 n=5+5) pmap_simple_2_devices 130µs ± 5% 20µs ± 5% -84.38% (p=0.008 n=5+5) pmap_simple_dispatch_8_devices 216µs ± 3% 39µs ± 5% -81.71% (p=0.008 n=5+5) pmap_simple_8_devices 221µs ± 6% 43µs ± 1% -80.41% (p=0.016 n=5+4) pmap_simple_dispatch_8_devices_100_args 5.52ms ± 7% 0.59ms ± 2% -89.28% (p=0.008 n=5+5) pmap_simple_8_devices_100_args 26.6ms ±21% 17.6ms ±37% -34.04% (p=0.008 n=5+5) sda_index_1 7.48µs ± 8% 7.53µs ± 6% ~ (p=0.310 n=5+5) sda_index_2 14.1µs ± 1% 14.3µs ± 4% ~ (p=0.690 n=5+5) sda_index_8 56.0µs ± 3% 56.9µs ± 4% ~ (p=0.310 n=5+5) ``` PiperOrigin-RevId: 401274089
2021-10-06 10:07:41 -07:00
* New features:
Add note to changelog about deprecation of jax.ops.index_...
2021-10-06 09:31:42 -04:00
* An optimized C++ code-path improving the dispatch time for `pmap` is now the
default when using jaxlib 0.1.72 or newer. The feature can be disabled using
the `--experimental_cpp_pmap` flag (or `JAX_CPP_PMAP` environment variable).
jnp.unique: add fill_value for when size is not None
2021-10-06 16:28:36 -07:00
* `jax.numpy.unique` now supports an optional `fill_value` argument ({jax-issue}`#8121`)
Create 0.2.21 jax release PiperOrigin-RevId: 398528427
2021-09-23 10:58:55 -07:00
Update CHANGELOG and verson numbers for jaxlib 0.1.72 release
2021-10-12 17:37:29 -07:00
## jaxlib 0.1.72 (Oct 12, 2021)
Remove 10.2 cuda support PiperOrigin-RevId: 402707900
2021-10-12 18:43:40 -07:00
* Breaking changes:
* Support for CUDA 10.2 and CUDA 10.1 has been dropped. Jaxlib now supports
CUDA 11.1+.
Add a regression test for #7461. Fixes #7461
2021-10-13 11:11:24 -04:00
* Bug fixes:
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* Fixes https://github.com/jax-ml/jax/issues/7461, which caused wrong
Add a regression test for #7461. Fixes #7461
2021-10-13 11:11:24 -04:00
outputs on all platforms due to incorrect buffer aliasing inside the XLA
compiler.
Update CHANGELOG and verson numbers for jaxlib 0.1.72 release
2021-10-12 17:37:29 -07:00
Create 0.2.21 jax release PiperOrigin-RevId: 398528427
2021-09-23 10:58:55 -07:00
## jax 0.2.21 (Sept 23, 2021)
* [GitHub
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
commits](https://github.com/jax-ml/jax/compare/jax-v0.2.20...jax-v0.2.21).
Add Changelog for jax v0.2.21 development
2021-09-15 12:10:30 -07:00
* Breaking Changes
Remove jax.api. Functions exported as jax.api were aliases for names in jax.*. Use the jax.* names instead.
2021-09-16 16:29:06 -04:00
* `jax.api` has been removed. Functions that were available as `jax.api.*`
were aliases for functions in `jax.*`; please use the functions in
`jax.*` instead.
Disallow non-hashable static arguments in pmap(). * Don't wrap static arguments in hashable wrappers in pmap. * Delete wrap_hashably(). * In argnums_partial, either enforce hashability or wrap values with an explicitly unhashable wrapper. The intent here is that either we should check for hashability early or we should make sure it's clear that it's not something we intended.. * Delete argnames_partial, which appears unused.
2021-07-19 13:11:38 -04:00
* `jax.partial`, and `jax.lax.partial` were accidental exports that have now
been removed. Use `functools.partial` from the Python standard library
instead.
jnp.array: raise TypeError on boolean scalar indices
2021-09-15 12:50:44 -07:00
* Boolean scalar indices now raise a `TypeError`; previously this silently
returned wrong results ({jax-issue}`#7925`).
Update changelog for several recent PRs
2021-09-16 14:10:08 -07:00
* Many more `jax.numpy` functions now require array-like inputs, and will error
if passed a list ({jax-issue}`#7747` {jax-issue}`#7802` {jax-issue}`#7907`).
See {jax-issue}`#7737` for a discussion of the rationale behind this change.
Attempt to land https://github.com/google/jax/pull/6400 again. This PR changes `jax.numpy.array()` to avoid creating any on-device arrays during tracing. As a consequence, calls to `jnp.array()` in a traced context, such as `jax.jit` will always be staged into the trace. This change may break code that depends on the current (undocumented and unintentional) behavior of `jnp.array()` to perform shape or index calculations that must be known statically (at trace time). The workaround for such cases is to use classic NumPy to perform shape/index calculations. PiperOrigin-RevId: 398008511
2021-09-21 09:05:56 -07:00
* When inside a transformation such as `jax.jit`, `jax.numpy.array` always
stages the array it produces into the traced computation. Previously
`jax.numpy.array` would sometimes produce a on-device array, even under
a `jax.jit` decorator. This change may break code that used JAX arrays to
perform shape or index computations that must be known statically; the
workaround is to perform such computations using classic NumPy arrays
instead.
Update v0.2.21 changelog for #7927
2021-09-27 11:38:36 -07:00
* `jnp.ndarray` is now a true base-class for JAX arrays. In particular, this
means that for a standard numpy array `x`, `isinstance(x, jnp.ndarray)` will
now return `False` ({jax-issue}`7927`).
Update changelog for several recent PRs
2021-09-16 14:10:08 -07:00
* New features:
* Added {func}`jax.numpy.insert` implementation ({jax-issue}`#7936`).
Add Changelog for jax v0.2.21 development
2021-09-15 12:10:30 -07:00
update version and changelog for pypi
2021-09-02 15:38:47 -07:00
## jax 0.2.20 (Sept 2, 2021)
update version and changelog for pypi
2021-08-12 21:17:53 -07:00
* [GitHub
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
commits](https://github.com/jax-ml/jax/compare/jax-v0.2.19...jax-v0.2.20).
update CHANGELOG for #7662 & #7732
2021-08-26 11:12:16 -07:00
* Breaking Changes
* `jnp.poly*` functions now require array-like inputs ({jax-issue}`#7732`)
* `jnp.unique` and other set-like operations now require array-like inputs
({jax-issue}`#7662`)
update version and changelog for pypi
2021-08-12 21:17:53 -07:00
Update files after new jaxlib release 0.1.71
2021-09-01 10:43:20 -07:00
## jaxlib 0.1.71 (Sep 1, 2021)
* Breaking changes:
* Support for CUDA 11.0 and CUDA 10.1 has been dropped. Jaxlib now supports
CUDA 10.2 and CUDA 11.1+.
update version and changelog for pypi
2021-08-12 21:17:53 -07:00
## jax 0.2.19 (Aug 12, 2021)
* [GitHub
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
commits](https://github.com/jax-ml/jax/compare/jax-v0.2.18...jax-v0.2.19).
Drop support for NumPy 1.17.
2021-07-29 09:18:01 -04:00
* Breaking changes:
* Support for NumPy 1.17 has been dropped, per the
[deprecation policy](https://jax.readthedocs.io/en/latest/deprecation.html).
Please upgrade to a supported NumPy version.
Add `@jit` decorators to jax.numpy operators. By wrapping common operators in `jit`, we get a number of benefits: * `jit` has a faster, more optimized dispatch path compared to the primitive dispatch path in JAX. It's faster to dispatch a `jit` computation than a single primitive. * `jit` allows us to cache and reuse logic such as broadcasting and type promotion. One downside is that we now report an error when large Python integer scalars (e.g. `2**32 - 1`) are passed as arguments to JAX array operators. The workaround to this is to use explicitly typed constants instead of Python scalars. On my laptop, this benchmark improves from 95us to 4us: ``` In [1]: import jax.numpy as jnp, jax In [2]: x = jax.device_put(7) WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.) In [3]: %timeit jnp.add(x, x).block_until_ready() 4.18 µs ± 159 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each) ``` PiperOrigin-RevId: 389871450
2021-08-10 06:48:55 -07:00
* The `jit` decorator has been added around the implementation of a number of
operators on JAX arrays. This speeds up dispatch times for common
operators such as `+`.
This change should largely be transparent to most users. However, there is
one known behavioral change, which is that large integer constants may now
produce an error when passed directly to a JAX operator
(e.g., `x + 2**40`). The workaround is to cast the constant to an
explicit type (e.g., `np.float64(2**40)`).
[jax2tf] Expand shape polymorphism support to use dimension polynomials as values. The goal of this change is to support shape polymorphism for operations such as average (which needs to divide by the size of a dimension) or indexing (which needs to normalize indices by comparing them with 0 and adding dimension size for negative indices). In both of these cases the size of a dimenion needs to be used as a value in the array computation. In general, the size of a dimension is used only to customize primitives. This change introduces `core.dim_as_value` which must be used on a dimension size before using it as a value in the array computation. E.g., ``` def average(x): return jnp.sum(x, axis=0) / core.dim_as_value(x.shape[0]) ``` This function is the identity function if the dimension size is constant, otherwise it uses a new primitive `shape_poly.dim_as_value_p`. Note that this does not change fundamentally the flavor of shape polymorphism supported in jax2tf: intermediate shapes and their values may depend on the input shapes, but never does a shape depend on the input values. In fact, one could have expressed the `dim_as_value` already: ``` def dim_as_value(d): jnp.sum(jnp.broadcast_to(jnp.array(1), shape=(d,))) ``` We were able to suppot `jnp.mean`, `jnp.average`, `jnp.take`, `lax.dynamic_slice`, `lax.dynamic_update_slice` by using `core.dim_as_value` internally, but to fully roll-up the solution we need to make `core.dim_as_value` a public API and teach the users how to use it when they want to use shape polymorphism. Alternatively, perhaps there is a way to automatically convert dimension polynomials to values when passed to the lax primitives.
2021-07-16 20:01:22 +03:00
* New features:
* Improved the support for shape polymorphism in jax2tf for operations that
need to use a dimension size in array computation, e.g., `jnp.mean`.
({jax-issue}`#7317`)
update version and changelog for pypi
2021-08-12 21:17:53 -07:00
* Bug fixes:
* Some leaked trace errors from the previous release ({jax-issue}`#7613`)
Upgrade versions after jaxlib release PiperOrigin-RevId: 389753047
2021-08-09 16:37:15 -07:00
## jaxlib 0.1.70 (Aug 9, 2021)
Bump jax version and CHANGELOG to 0.2.18
2021-07-21 11:56:24 -07:00
* Breaking changes:
* Support for Python 3.6 has been dropped, per the
[deprecation policy](https://jax.readthedocs.io/en/latest/deprecation.html).
Please upgrade to a supported Python version.
Drop support for NumPy 1.17.
2021-07-29 09:18:01 -04:00
* Support for NumPy 1.17 has been dropped, per the
[deprecation policy](https://jax.readthedocs.io/en/latest/deprecation.html).
Please upgrade to a supported NumPy version.
[jax2tf] Expand shape polymorphism support to use dimension polynomials as values. The goal of this change is to support shape polymorphism for operations such as average (which needs to divide by the size of a dimension) or indexing (which needs to normalize indices by comparing them with 0 and adding dimension size for negative indices). In both of these cases the size of a dimenion needs to be used as a value in the array computation. In general, the size of a dimension is used only to customize primitives. This change introduces `core.dim_as_value` which must be used on a dimension size before using it as a value in the array computation. E.g., ``` def average(x): return jnp.sum(x, axis=0) / core.dim_as_value(x.shape[0]) ``` This function is the identity function if the dimension size is constant, otherwise it uses a new primitive `shape_poly.dim_as_value_p`. Note that this does not change fundamentally the flavor of shape polymorphism supported in jax2tf: intermediate shapes and their values may depend on the input shapes, but never does a shape depend on the input values. In fact, one could have expressed the `dim_as_value` already: ``` def dim_as_value(d): jnp.sum(jnp.broadcast_to(jnp.array(1), shape=(d,))) ``` We were able to suppot `jnp.mean`, `jnp.average`, `jnp.take`, `lax.dynamic_slice`, `lax.dynamic_update_slice` by using `core.dim_as_value` internally, but to fully roll-up the solution we need to make `core.dim_as_value` a public API and teach the users how to use it when they want to use shape polymorphism. Alternatively, perhaps there is a way to automatically convert dimension polynomials to values when passed to the lax primitives.
2021-07-16 20:01:22 +03:00
* The host_callback mechanism now uses one thread per local device for
Bump jax version and CHANGELOG to 0.2.18
2021-07-21 11:56:24 -07:00
making the calls to the Python callbacks. Previously there was a single
thread for all devices. This means that the callbacks may now be called
interleaved. The callbacks corresponding to one device will still be
called in sequence.
## jax 0.2.18 (July 21 2021)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.2.17...jax-v0.2.18).
Update jax version and changelog for 0.1.27. Disable tfrt CPU backend on jaxlib 0.1.68 to work around https://github.com/google/jax/issues/7229.
2021-07-09 15:19:24 -04:00
Drop Python 3.6 support. Per the deprecation policy (https://jax.readthedocs.io/en/latest/deprecation.html), Python 3.6 support has been due for removal since June 23, 2020.
2021-07-15 09:50:30 -04:00
* Breaking changes:
* Support for Python 3.6 has been dropped, per the
[deprecation policy](https://jax.readthedocs.io/en/latest/deprecation.html).
Please upgrade to a supported Python version.
Update minimum jaxlib version to 0.1.69.
2021-07-15 16:39:18 -04:00
* The minimum jaxlib version is now 0.1.69.
* The `backend` argument to {py:func}`jax.dlpack.from_dlpack` has been
removed.
Drop Python 3.6 support. Per the deprecation policy (https://jax.readthedocs.io/en/latest/deprecation.html), Python 3.6 support has been due for removal since June 23, 2020.
2021-07-15 09:50:30 -04:00
Remove additional info return value from jax.scipy.linalg.polar().
2021-07-20 13:12:13 -04:00
* New features:
* Added a polar decomposition ({py:func}`jax.scipy.linalg.polar`).
[jax2tf] Fix conversion for argmin/argmax; add conversion for reduce The previous conversion for argmin/argmax simply used tf.argmin and tf.argmax. Those ops behave differently than JAX when the inputs contain NaN and Inf. Added a few test cases in primitive_harness to expose the failures. In order to implement an accurate conversion of argmin/argmax, we need to use the XLA Reduce op. Also tightened the shape checks for lax.argmin and lax.argmax, to ensure they are not used with an empty reduced dimension. E.g., if the axis=-1, previously we got an internal error: ``` RuntimeError: Invalid argument: Reducing out-of-bounds dimension -1 in shape f32[2,0,3].: This is a bug in JAX's shape-checking rules; please report it! ``` PiperOrigin-RevId: 384182794
2021-07-12 01:11:17 -07:00
* Bug fixes:
* Tightened the checks for lax.argmin and lax.argmax to ensure they are
DOC: many small fixes
2021-08-02 17:57:09 -07:00
not used with an invalid `axis` value, or with an empty reduction dimension.
[jax2tf] Fix conversion for argmin/argmax; add conversion for reduce The previous conversion for argmin/argmax simply used tf.argmin and tf.argmax. Those ops behave differently than JAX when the inputs contain NaN and Inf. Added a few test cases in primitive_harness to expose the failures. In order to implement an accurate conversion of argmin/argmax, we need to use the XLA Reduce op. Also tightened the shape checks for lax.argmin and lax.argmax, to ensure they are not used with an empty reduced dimension. E.g., if the axis=-1, previously we got an internal error: ``` RuntimeError: Invalid argument: Reducing out-of-bounds dimension -1 in shape f32[2,0,3].: This is a bug in JAX's shape-checking rules; please report it! ``` PiperOrigin-RevId: 384182794
2021-07-12 01:11:17 -07:00
({jax-issue}`#7196`)
Update minimum jaxlib version to 0.1.69.
2021-07-15 16:39:18 -04:00
Update changelog for jaxlib-0.1.69.
2021-07-09 21:16:48 -07:00
## jaxlib 0.1.69 (July 9 2021)
* Fix bugs in TFRT CPU backend that results in incorrect results.
Update jax version and changelog for 0.1.27. Disable tfrt CPU backend on jaxlib 0.1.68 to work around https://github.com/google/jax/issues/7229.
2021-07-09 15:19:24 -04:00
## jax 0.2.17 (July 9 2021)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.2.16...jax-v0.2.17).
Update jax version and changelog for 0.1.27. Disable tfrt CPU backend on jaxlib 0.1.68 to work around https://github.com/google/jax/issues/7229.
2021-07-09 15:19:24 -04:00
* Bug fixes:
* Default to the older "stream_executor" CPU runtime for jaxlib <= 0.1.68
to work around #7229, which caused wrong outputs on CPU due to a concurrency
problem.
Adds spherical harmonics. Co-authored-by: Jake VanderPlas <jakevdp@google.com>
2021-07-02 10:42:29 -07:00
* New features:
* New SciPy function {py:func}`jax.scipy.special.sph_harm`.
AWN-enabled reduction over named axes in reverse-mode AD Previously, reverse-mode AD operators inside JAX maps always meant "compute a gradient (or VJP, etc.) for each axis index in the map". For instance, `vmap(grad(f))` is the standard JAX spelling of the per-example gradient of `f`. In batching tracer terms, this "elementwise" behavior means that, if any inputs to a function being transposed are mapped, the cotangents of all inputs, even unmapped ones, would also be mapped. But a user might want them to be unmapped (if, for instance, they're interested in a total gradient rather than a per-example gradient). They could always reduce (`psum`) the cotangents afterwards, but computing mapped cotangents in the first place would likely be an unacceptable waste of memory and can't necessarily be optimized away. If we want to fuse these reductions into reverse-mode autodiff itself, we need the backward_pass logic and/or transpose rules to know about whether primal values are mapped or unmapped. This is made possible by avals-with-names, which encodes that information in the avals of the primal jaxpr. Putting things together, **this change adds an option to reverse-mode AD APIs that indicates which named axes should be reduced over in the backward pass in situations where they were broadcasted over in the forward pass**. All other named axes will be treated in the current elementwise way. This has the effect of making APIs like `grad` behave akin to collectives like `psum`: they act collectively over axes that are named explicitly, and elementwise otherwise. Since avals-with-names is currently enabled only in `xmap`, this behavior is only available in that context for now. It's also missing some optimizations: - reductions aren't fused into any first-order primitives (e.g. a `pdot` should have a named contracting axis added rather than being followed by a `psum`; this can be implemented by putting these primitives into `reducing_transposes`) - reductions are performed eagerly, even over axes that are mapped to hardware resources (the optimal thing to do would be to reduce eagerly over any vectorized axis component while delaying the reduction over any hardware-mapped component until the end of the overall backward pass; this would require a way to represent these partially-reduced values) PiperOrigin-RevId: 383685336
2021-07-08 12:05:56 -07:00
* Reverse-mode autodiff functions ({func}`jax.grad`,
{func}`jax.value_and_grad`, {func}`jax.vjp`, and
{func}`jax.linear_transpose`) support a parameter that indicates which named
axes should be summed over in the backward pass if they were broadcasted
over in the forward pass. This enables use of these APIs in a
non-per-example way inside maps (initially only
{func}`jax.experimental.maps.xmap`) ({jax-issue}`#6950`).
Update jax version to 0.2.16
2021-06-23 14:42:36 -07:00
Update changelog for jax and jaxlib releases.
2021-06-28 13:52:19 -07:00
Update jax version to 0.2.16
2021-06-23 14:42:36 -07:00
## jax 0.2.16 (June 23 2021)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.2.15...jax-v0.2.16).
Update jax version to 0.2.15
2021-06-23 11:55:40 -07:00
## jax 0.2.15 (June 23 2021)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.2.14...jax-v0.2.15).
[jax2tf] Fix the round-trip call_tf(convert) Also cleaned the handling of global state in jax2tf.
2021-06-10 17:01:22 +02:00
* New features:
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [#7042](https://github.com/jax-ml/jax/pull/7042) Turned on TFRT CPU backend
Update changelog for jax and jaxlib releases.
2021-06-28 13:52:19 -07:00
with significant dispatch performance improvements on CPU.
[jax2tf] Support inequality and min/max for booleans. For inequalities we add casts to int8. For min/max we rewrite to logical operations and/or.
2021-06-10 12:42:40 +02:00
* The {func}`jax2tf.convert` supports inequalities and min/max for booleans
({jax-issue}`#6956`).
Adds spherical harmonics. Co-authored-by: Jake VanderPlas <jakevdp@google.com>
2021-07-02 10:42:29 -07:00
* New SciPy function {py:func}`jax.scipy.special.lpmn_values`.
Update jax version to 0.2.14
2021-06-10 13:15:53 -07:00
[jax2tf] Fix the round-trip call_tf(convert) Also cleaned the handling of global state in jax2tf.
2021-06-10 17:01:22 +02:00
* Breaking changes:
Drop support for NumPy 1.16.
2021-06-10 12:12:13 -04:00
* Support for NumPy 1.16 has been dropped, per the
[deprecation policy](https://jax.readthedocs.io/en/latest/deprecation.html).
Update jax version to 0.2.14
2021-06-10 13:15:53 -07:00
[jax2tf] Fix the round-trip call_tf(convert) Also cleaned the handling of global state in jax2tf.
2021-06-10 17:01:22 +02:00
* Bug fixes:
* Fixed bug that prevented round-tripping from JAX to TF and back:
`jax2tf.call_tf(jax2tf.convert)` ({jax-issue}`#6947`).
Update jax version to 0.2.15
2021-06-23 11:55:40 -07:00
Update changelog for jax and jaxlib releases.
2021-06-28 13:52:19 -07:00
## jaxlib 0.1.68 (June 23 2021)
* Bug fixes:
* Fixed bug in TFRT CPU backend that gets nans when transfer TPU buffer to
CPU.
Update jax version to 0.2.14
2021-06-10 13:15:53 -07:00
## jax 0.2.14 (June 10 2021)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.2.13...jax-v0.2.14).
[jax2tf] Fix the scoping of the enable_xla conversion parameter Previously, the global enable_xla flag was set upon entry to `jax.convert`. It should instead be set only for the duration of the just-in-time conversion, which may happen later when the converted function is invoked.
2021-05-11 11:11:37 +03:00
* New features:
Merge pull request #6866 from gnecula:tf_pjit PiperOrigin-RevId: 376989780
2021-06-01 22:50:12 -07:00
* The {func}`jax2tf.convert` now has support for `pjit` and `sharded_jit`.
Add a --jax_traceback_filtering flag to control the traceback filtering mode. Add a new traceback filtering mode that uses __tracebackhide__, and use it in IPython.
2021-06-02 15:22:50 -04:00
* A new configuration option JAX_TRACEBACK_FILTERING controls how JAX filters
tracebacks.
* A new traceback filtering mode using `__tracebackhide__` is now enabled by
default in sufficiently recent versions of IPython.
[jax2tf] Implement inequalities and friends for complex numbers. This requires re-using JAX's lowering rule for comparisons of complex numbers to use lexicographic comparison.
2021-06-04 11:02:50 +03:00
* The {func}`jax2tf.convert` supports shape polymorphism even when the
unknown dimensions are used in arithmetic operations, e.g., `jnp.reshape(-1)`
Update jax version to 0.2.14
2021-06-10 13:15:53 -07:00
({jax-issue}`#6827`).
[jax2tf] Add support for generating HLO OpMetadata in the TF graph The goal is to ensure that the HLO that jax2tf->TF/XLA generates has the same metadata as what JAX generates. This includes `op_type`, `op_name`, and source information, which are used for debugging and profiling. In order to ensure that this metadata is carried from the JAX tracing time to TF/XLA, we save the metadata in custom TF op attributes. These attributes are automatically preserved through SavedModel. This relies on a separate change in TF/XLA to look for these custom attributes and override its default. For the source information, we use pretty much the same code that xla.py uses. HLO OpMetadata has room for only one source location. JAX (xla.py) picks the top-most user frame, which is obtained by filtering out the stack frames in the JAX source tree. When used with jax2tf we also need to filter out stack frames in the TensorFlow source tree. The hardest part is to generate the `op_name`, which is a hierarchical name with components separated by '/', e.g., `jax2tf(top_func)/while/cond/le`. We carry the current `name_stack` in thread-local state. Unfortunately, there is no easy way to share the exact code that achieves this in xla.py. At the same time it is not crucial that we have exactly identical name stacks as in JAX. I attempted to also carry this state in the JAX `MainTrace`, but could not fully control the name stack. E.g., when calling a jitted-function we have to reuse the current `MainTrace` although we want to push an element on the name stack. For now this option is not yet enabled until we make the necessary changes in TensorFlow.
2021-05-25 13:33:35 +02:00
* The {func}`jax2tf.convert` generates custom attributes with location information
in TF ops. The code that XLA generates after jax2tf
has the same location information as JAX/XLA.
Adds spherical harmonics. Co-authored-by: Jake VanderPlas <jakevdp@google.com>
2021-07-02 10:42:29 -07:00
* New SciPy function {py:func}`jax.scipy.special.lpmn`.
[jax2tf] Add support for generating HLO OpMetadata in the TF graph The goal is to ensure that the HLO that jax2tf->TF/XLA generates has the same metadata as what JAX generates. This includes `op_type`, `op_name`, and source information, which are used for debugging and profiling. In order to ensure that this metadata is carried from the JAX tracing time to TF/XLA, we save the metadata in custom TF op attributes. These attributes are automatically preserved through SavedModel. This relies on a separate change in TF/XLA to look for these custom attributes and override its default. For the source information, we use pretty much the same code that xla.py uses. HLO OpMetadata has room for only one source location. JAX (xla.py) picks the top-most user frame, which is obtained by filtering out the stack frames in the JAX source tree. When used with jax2tf we also need to filter out stack frames in the TensorFlow source tree. The hardest part is to generate the `op_name`, which is a hierarchical name with components separated by '/', e.g., `jax2tf(top_func)/while/cond/le`. We carry the current `name_stack` in thread-local state. Unfortunately, there is no easy way to share the exact code that achieves this in xla.py. At the same time it is not crucial that we have exactly identical name stacks as in JAX. I attempted to also carry this state in the JAX `MainTrace`, but could not fully control the name stack. E.g., when calling a jitted-function we have to reuse the current `MainTrace` although we want to push an element on the name stack. For now this option is not yet enabled until we make the necessary changes in TensorFlow.
2021-05-25 13:33:35 +02:00
[jax2tf] Fix the scoping of the enable_xla conversion parameter Previously, the global enable_xla flag was set upon entry to `jax.convert`. It should instead be set only for the duration of the just-in-time conversion, which may happen later when the converted function is invoked.
2021-05-11 11:11:37 +03:00
* Bug fixes:
[jax2tf] Fix the 32/64-bit behavior to follow JAX rules JAX and TensorFlow have different behavior w.r.t. 32-64 bit computations. This PR cleans up the handling of types in jax2tf to ensure that we follow the same behavior in jax2tf and in JAX. This means that f_jax(args) always does the computation with the same precision as jax2tf.convert(f_jax)(args). This may mean that the result of the conversion depends on the value of JAX_ENABLE_x64. See README.md for more details.
2021-06-02 21:03:16 +03:00
* The {func}`jax2tf.convert` now ensures that it uses the same typing rules
for Python scalars and for choosing 32-bit vs. 64-bit computations
as JAX ({jax-issue}`#6883`).
[jax2tf] Fix the scoping of the enable_xla conversion parameter Previously, the global enable_xla flag was set upon entry to `jax.convert`. It should instead be set only for the duration of the just-in-time conversion, which may happen later when the converted function is invoked.
2021-05-11 11:11:37 +03:00
* The {func}`jax2tf.convert` now scopes the `enable_xla` conversion parameter
properly to apply only during the just-in-time conversion
({jax-issue}`#6720`).
Copybara import of the project: -- 1ecf4f02891cad70cc8f094b49cf2458105ca366 by George Necula <gcnecula@gmail.com>: [jax2tf] Change the conversion of dot_general to use XLA op. Instead of converting the dot_general to a sea of TF ops, when we enable_xla we just use the XLA op. This has the advantage that it also supports the preferred_element_type. Fixed bug with passing the precision parameter to TF. Also improved tests to print the HLO in case of numerical errors. COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/6717 from gnecula:tf_dot 1ecf4f02891cad70cc8f094b49cf2458105ca366 PiperOrigin-RevId: 373326655
2021-05-12 02:29:51 -07:00
* The {func}`jax2tf.convert` now converts `lax.dot_general` using the
`XlaDot` TensorFlow op, for better fidelity w.r.t. JAX numerical precision
({jax-issue}`#6717`).
[jax2tf] Implement inequalities and friends for complex numbers. This requires re-using JAX's lowering rule for comparisons of complex numbers to use lexicographic comparison.
2021-06-04 11:02:50 +03:00
* The {func}`jax2tf.convert` now has support for inequality comparisons and
Update jax version to 0.2.14
2021-06-10 13:15:53 -07:00
min/max for complex numbers ({jax-issue}`#6892`).
Copybara import of the project: -- 1ecf4f02891cad70cc8f094b49cf2458105ca366 by George Necula <gcnecula@gmail.com>: [jax2tf] Change the conversion of dot_general to use XLA op. Instead of converting the dot_general to a sea of TF ops, when we enable_xla we just use the XLA op. This has the advantage that it also supports the preferred_element_type. Fixed bug with passing the precision parameter to TF. Also improved tests to print the HLO in case of numerical errors. COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/6717 from gnecula:tf_dot 1ecf4f02891cad70cc8f094b49cf2458105ca366 PiperOrigin-RevId: 373326655
2021-05-12 02:29:51 -07:00
Update changelog for jax and jaxlib releases.
2021-06-28 13:52:19 -07:00
## jaxlib 0.1.67 (May 17 2021)
Update README etc for jaxlib 0.1.66 release.
2021-05-11 16:49:32 -07:00
## jaxlib 0.1.66 (May 11 2021)
Copybara import of the project: -- 1ecf4f02891cad70cc8f094b49cf2458105ca366 by George Necula <gcnecula@gmail.com>: [jax2tf] Change the conversion of dot_general to use XLA op. Instead of converting the dot_general to a sea of TF ops, when we enable_xla we just use the XLA op. This has the advantage that it also supports the preferred_element_type. Fixed bug with passing the precision parameter to TF. Also improved tests to print the HLO in case of numerical errors. COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/6717 from gnecula:tf_dot 1ecf4f02891cad70cc8f094b49cf2458105ca366 PiperOrigin-RevId: 373326655
2021-05-12 02:29:51 -07:00
Bump jax version and changelog for jax 0.1.13 release
2021-05-03 16:32:00 -07:00
* New features:
* CUDA 11.1 wheels are now supported on all CUDA 11 versions 11.1 or higher.
NVidia now promises compatibility between CUDA minor releases starting with
CUDA 11.1. This means that JAX can release a single CUDA 11.1 wheel that
is compatible with CUDA 11.2 and 11.3.
There is no longer a separate jaxlib release for CUDA 11.2 (or higher); use
the CUDA 11.1 wheel for those versions (cuda111).
Update changelog.
2021-05-04 09:28:04 -04:00
* Jaxlib now bundles `libdevice.10.bc` in CUDA wheels. There should be no need
to point JAX to a CUDA installation to find this file.
* Added automatic support for static keyword arguments to the {func}`jit`
implementation.
* Added support for pretransformation exception traces.
* Initial support for pruning unused arguments from {func}`jit` -transformed
computations.
Pruning is still a work in progress.
* Improved the string representation of {class}`PyTreeDef` objects.
* Added support for XLA's variadic ReduceWindow.
* Bug fixes:
* Fixed a bug in the remote cloud TPU support when large numbers of arguments
are passed to a computation.
* Fix a bug that meant that JAX garbage collection was not triggered by
{func}`jit` transformed functions.
Bump jax version and changelog for jax 0.1.13 release
2021-05-03 16:32:00 -07:00
## jax 0.2.13 (May 3 2021)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.2.12...jax-v0.2.13).
[JAX] Add static_argnames support to jax.jit. Requires a new jaxlib build. Add support for static_argnames in C++ JIT implementation. PiperOrigin-RevId: 367627359
2021-04-09 07:10:02 -07:00
* New features:
* When combined with jaxlib 0.1.66, {func}`jax.jit` now supports static
keyword arguments. A new `static_argnames` option has been added to specify
keyword arguments as static.
Add JIT-compatible version of jnp.nonzero
2021-04-20 09:18:26 -07:00
* {func}`jax.nonzero` has a new optional `size` argument that allows it to
jnp.unique: add support for axis argument
2021-04-21 16:00:14 -07:00
be used within `jit` ({jax-issue}`#6501`)
* {func}`jax.numpy.unique` now supports the `axis` argument ({jax-issue}`#6532`).
[host_callback] Minor fix to use the new xla_shape.is_token
2021-04-28 12:22:32 +03:00
* {func}`jax.experimental.host_callback.call` now supports `pjit.pjit` ({jax-issue}`#6569`).
Add a tridiagonal eigh solver.
2021-05-03 11:27:07 -04:00
* Added {func}`jax.scipy.linalg.eigh_tridiagonal` that computes the
eigenvalues of a tridiagonal matrix. Only eigenvalues are supported at
present.
Update changelog.
2021-05-04 09:28:04 -04:00
* The order of the filtered and unfiltered stack traces in exceptions has been
changed. The traceback attached to an exception thrown from JAX-transformed
code is now filtered, with an `UnfilteredStackTrace` exception
containing the original trace as the `__cause__` of the filtered exception.
Filtered stack traces now also work with Python 3.6.
* If an exception is thrown by code that has been transformed by reverse-mode
automatic differentiation, JAX now attempts to attach as a `__cause__` of
the exception a `JaxStackTraceBeforeTransformation` object that contains the
stack trace that created the original operation in the forward pass.
Requires jaxlib 0.1.66.
[JAX] Add static_argnames support to jax.jit. Requires a new jaxlib build. Add support for static_argnames in C++ JIT implementation. PiperOrigin-RevId: 367627359
2021-04-09 07:10:02 -07:00
* Breaking changes:
Replace `host_id` with `process_index` terminology, take 2. We're switching to the new terminology to avoid confusion in cases where multiple jax processes are running on a single host, and each process has a unique process_index/host_id. This keeps aliases for the old `host_id` APIs for now, but these will eventually be removed. This was originally commited in b77ef5138b631378e6a8ceb8bafc94fe91239bae, but reverted in 14acd070c2afb11c81fc91f43790577cd48cbf67 due to Google-internal test failures from renaming the local_devices argument name. This change is identical except it also adds staging for the argument name change.
2021-04-20 17:56:41 -07:00
* The following function names have changed. There are still aliases, so this
should not break existing code, but the aliases will eventually be removed
so please change your code.
* `host_id` --> {func}`~jax.process_index`
* `host_count` --> {func}`~jax.process_count`
* `host_ids` --> `range(jax.process_count())`
* Similarly, the argument to {func}`~jax.local_devices` has been renamed from
`host_id` to `process_index`.
Mark arguments to jax.jit() other than the function as keyword-only. This change is to prevent breakage when options are added or removed.
2021-04-08 10:32:35 -04:00
* Arguments to {func}`jax.jit` other than the function are now marked as
keyword-only. This change is to prevent accidental breakage when arguments
are added to `jit`.
[JAX] Add static_argnames support to jax.jit. Requires a new jaxlib build. Add support for static_argnames in C++ JIT implementation. PiperOrigin-RevId: 367627359
2021-04-09 07:10:02 -07:00
* Bug fixes:
* The {func}`jax2tf.convert` now works in presence of gradients for functions
with integer inputs ({jax-issue}`#6360`).
[host_callback] Minor fix to use the new xla_shape.is_token
2021-04-28 12:22:32 +03:00
* Fixed assertion failure in {func}`jax2tf.call_tf` when used with captured
`tf.Variable` ({jax-issue}`#6572`).
Update jax version and changelog
2021-04-01 10:11:52 -07:00
Update README, etc. for jaxlib 0.1.65 release
2021-04-07 17:51:20 -07:00
## jaxlib 0.1.65 (April 7 2021)
Update jax version and changelog
2021-04-01 10:11:52 -07:00
## jax 0.2.12 (April 1 2021)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.2.11...v0.2.12).
Add programmatic profiling APIs, and rename some existing APIs. This change provides aliases for the renamed APIs so existing code won't break. We should remove these aliases after the next release.
2021-03-17 17:33:25 +00:00
* New features
* New profiling APIs: {func}`jax.profiler.start_trace`,
{func}`jax.profiler.stop_trace`, and {func}`jax.profiler.trace`
Add a JVP rule for the general case of lax.reduce.
2021-03-23 10:31:02 -04:00
* {func}`jax.lax.reduce` is now differentiable.
Increase minimum jaxlib version to 0.1.64.
2021-03-24 13:49:05 -04:00
* Breaking changes:
* The minimum jaxlib version is now 0.1.64.
Add programmatic profiling APIs, and rename some existing APIs. This change provides aliases for the renamed APIs so existing code won't break. We should remove these aliases after the next release.
2021-03-17 17:33:25 +00:00
* Some profiler APIs names have been changed. There are still aliases, so this
should not break existing code, but the aliases will eventually be removed
so please change your code.
* `TraceContext` --> {func}`~jax.profiler.TraceAnnotation`
* `StepTraceContext` --> {func}`~jax.profiler.StepTraceAnnotation`
* `trace_function` --> {func}`~jax.profiler.annotate_function`
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* Omnistaging can no longer be disabled. See [omnistaging](https://github.com/jax-ml/jax/blob/main/docs/design_notes/omnistaging.md)
pre-release omnistaging cleanup
2021-03-25 16:44:58 -07:00
for more information.
Python integer conversion: always return int64 or OverflowError
2021-03-29 09:26:19 -07:00
* Python integers larger than the maximum `int64` value will now lead to an overflow
in all cases, rather than being silently converted to `uint64` in some cases ({jax-issue}`#6047`).
X32 mode: raise OverflowError for large integers
2021-03-30 10:05:03 -07:00
* Outside X64 mode, Python integers outside the range representable by `int32` will now lead to an
`OverflowError` rather than having their value silently truncated.
[host_callback] Add support for tapping empty arrays We make sure that both the inputs and the outputs of callbacks can contain empty arrays. Most platforms do not support empty infeed, so we ensure we do not send those.
2021-03-29 17:21:56 +03:00
* Bug fixes:
* `host_callback` now supports empty arrays in arguments and results ({jax-issue}`#6262`).
Update changelog for #5868
2021-04-01 09:29:22 -07:00
* {func}`jax.random.randint` clips rather than wraps of out-of-bounds limits, and can now generate
integers in the full range of the specified dtype ({jax-issue}`#5868`)
Increase minimum jaxlib version to 0.1.64.
2021-03-24 13:49:05 -04:00
update version and changelog for pypi
2021-03-23 19:33:04 -07:00
## jax 0.2.11 (March 23 2021)
* [GitHub
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
commits](https://github.com/jax-ml/jax/compare/jax-v0.2.10...jax-v0.2.11).
move changelog to top level
2021-03-05 11:07:50 -08:00
* New features:
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [#6112](https://github.com/jax-ml/jax/pull/6112) added context managers:
update version and changelog for pypi
2021-03-23 19:33:04 -07:00
`jax.enable_checks`, `jax.check_tracer_leaks`, `jax.debug_nans`,
`jax.debug_infs`, `jax.log_compiles`.
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [#6085](https://github.com/jax-ml/jax/pull/6085) added `jnp.delete`
Add a JVP rule for the general case of lax.reduce.
2021-03-23 10:31:02 -04:00
move changelog to top level
2021-03-05 11:07:50 -08:00
* Bug fixes:
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [#6136](https://github.com/jax-ml/jax/pull/6136) generalized
update version and changelog for pypi
2021-03-23 19:33:04 -07:00
`jax.flatten_util.ravel_pytree` to handle integer dtypes.
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [#6129](https://github.com/jax-ml/jax/issues/6129) fixed a bug with handling
update version and changelog for pypi
2021-03-23 19:33:04 -07:00
some constants like `enum.IntEnums`
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [#6145](https://github.com/jax-ml/jax/pull/6145) fixed batching issues with
update version and changelog for pypi
2021-03-23 19:33:04 -07:00
incomplete beta functions
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [#6014](https://github.com/jax-ml/jax/pull/6014) fixed H2D transfers during
update version and changelog for pypi
2021-03-23 19:33:04 -07:00
tracing
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [#6165](https://github.com/jax-ml/jax/pull/6165) avoids OverflowErrors when
update version and changelog for pypi
2021-03-23 19:33:04 -07:00
converting some large Python integers to floats
move changelog to top level
2021-03-05 11:07:50 -08:00
* Breaking changes:
Increase minimum jaxlib version to 0.1.62.
2021-03-16 12:13:41 -04:00
* The minimum jaxlib version is now 0.1.62.
move changelog to top level
2021-03-05 11:07:50 -08:00
Update README, etc. for jaxlib 0.1.64 release
2021-03-18 16:11:40 -07:00
## jaxlib 0.1.64 (March 18 2021)
Update README, etc. for jaxlib 0.1.63 release
2021-03-17 10:14:52 -07:00
## jaxlib 0.1.63 (March 17 2021)
Update jaxlib version etc.
2021-03-09 17:55:40 -08:00
move changelog to top level
2021-03-05 11:07:50 -08:00
## jax 0.2.10 (March 5 2021)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.2.9...jax-v0.2.10).
move changelog to top level
2021-03-05 11:07:50 -08:00
* New features:
* {func}`jax.scipy.stats.chi2` is now available as a distribution with logpdf and pdf methods.
* {func}`jax.scipy.stats.betabinom` is now available as a distribution with logpmf and pmf methods.
* Added {func}`jax.experimental.jax2tf.call_tf` to call TensorFlow functions
from JAX ({jax-issue}`#5627`)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
and [README](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/README.md#calling-tensorflow-functions-from-jax)).
move changelog to top level
2021-03-05 11:07:50 -08:00
* Extended the batching rule for `lax.pad` to support batching of the padding values.
* Bug fixes:
* {func}`jax.numpy.take` properly handles negative indices ({jax-issue}`#5768`)
* Breaking changes:
* JAX's promotion rules were adjusted to make promotion more consistent and
invariant to JIT. In particular, binary operations can now result in weakly-typed
values when appropriate. The main user-visible effect of the change is that
some operations result in outputs of different precision than before; for
example the expression `jnp.bfloat16(1) + 0.1 * jnp.arange(10)`
previously returned a `float64` array, and now returns a `bfloat16` array.
JAX's type promotion behavior is described at {ref}`type-promotion`.
* {func}`jax.numpy.linspace` now computes the floor of integer values, i.e.,
rounding towards -inf rather than 0. This change was made to match NumPy
1.20.0.
* {func}`jax.numpy.i0` no longer accepts complex numbers. Previously the
function computed the absolute value of complex arguments. This change was
made to match the semantics of NumPy 1.20.0.
* Several {mod}`jax.numpy` functions no longer accept tuples or lists in place
of array arguments: {func}`jax.numpy.pad`, :func`jax.numpy.ravel`,
{func}`jax.numpy.repeat`, {func}`jax.numpy.reshape`.
In general, {mod}`jax.numpy` functions should be used with scalars or array arguments.
Update jaxlib version etc.
2021-03-09 17:55:40 -08:00
## jaxlib 0.1.62 (March 9 2021)
move changelog to top level
2021-03-05 11:07:50 -08:00
* New features:
* jaxlib wheels are now built to require AVX instructions on x86-64 machines
by default. If you want to use JAX on a machine that doesn't support AVX,
you can build a jaxlib from source using the `--target_cpu_features` flag
to `build.py`. `--target_cpu_features` also replaces
`--enable_march_native`.
## jaxlib 0.1.61 (February 12 2021)
Fix Typos in CHANGELOG.md This PR fixes the typos in Change log documentation
2024-03-12 13:57:07 +05:30
## jaxlib 0.1.60 (February 3 2021)
move changelog to top level
2021-03-05 11:07:50 -08:00
* Bug fixes:
* Fixed a memory leak when converting CPU DeviceArrays to NumPy arrays. The
memory leak was present in jaxlib releases 0.1.58 and 0.1.59.
* `bool`, `int8`, and `uint8` are now considered safe to cast to
`bfloat16` NumPy extension type.
## jax 0.2.9 (January 26 2021)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.2.8...jax-v0.2.9).
move changelog to top level
2021-03-05 11:07:50 -08:00
* New features:
* Extend the {mod}`jax.experimental.loops` module with support for pytrees. Improved
error checking and error messages.
* Add {func}`jax.experimental.enable_x64` and {func}`jax.experimental.disable_x64`.
These are context managers which allow X64 mode to be temporarily enabled/disabled
within a session.
* Breaking changes:
* {func}`jax.ops.segment_sum` now drops segment IDs that are out of range rather
than wrapping them into the segment ID space. This was done for performance
reasons.
## jaxlib 0.1.59 (January 15 2021)
## jax 0.2.8 (January 12 2021)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.2.7...jax-v0.2.8).
move changelog to top level
2021-03-05 11:07:50 -08:00
* New features:
* Add {func}`jax.closure_convert` for use with higher-order custom
derivative functions. ({jax-issue}`#5244`)
* Add {func}`jax.experimental.host_callback.call` to call a custom Python
function on the host and return a result to the device computation.
({jax-issue}`#5243`)
* Bug fixes:
* `jax.numpy.arccosh` now returns the same branch as `numpy.arccosh` for
complex inputs ({jax-issue}`#5156`)
DOC: many small fixes
2021-08-02 17:57:09 -07:00
* `host_callback.id_tap` now works for `jax.pmap` also. There is an
move changelog to top level
2021-03-05 11:07:50 -08:00
optional parameter for `id_tap` and `id_print` to request that the
device from which the value is tapped be passed as a keyword argument
to the tap function ({jax-issue}`#5182`).
* Breaking changes:
* `jax.numpy.pad` now takes keyword arguments. Positional argument `constant_values`
has been removed. In addition, passing unsupported keyword arguments raises an error.
* Changes for {func}`jax.experimental.host_callback.id_tap` ({jax-issue}`#5243`):
* Removed support for `kwargs` for {func}`jax.experimental.host_callback.id_tap`.
(This support has been deprecated for a few months.)
* Changed the printing of tuples for {func}`jax.experimental.host_callback.id_print`
to use '(' instead of '['.
* Changed the {func}`jax.experimental.host_callback.id_print` in presence of JVP
to print a pair of primal and tangent. Previously, there were two separate
print operations for the primals and the tangent.
* `host_callback.outfeed_receiver` has been removed (it is not necessary,
and was deprecated a few months ago).
* New features:
Fix Typos in CHANGELOG.md This PR fixes the typos in Change log documentation
2024-03-12 13:57:07 +05:30
* New flag for debugging `inf`, analogous to that for `NaN` ({jax-issue}`#5224`).
move changelog to top level
2021-03-05 11:07:50 -08:00
## jax 0.2.7 (Dec 4 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.2.6...jax-v0.2.7).
move changelog to top level
2021-03-05 11:07:50 -08:00
* New features:
* Add `jax.device_put_replicated`
* Add multi-host support to `jax.experimental.sharded_jit`
DOC: many small fixes
2021-08-02 17:57:09 -07:00
* Add support for differentiating eigenvalues computed by `jax.numpy.linalg.eig`
move changelog to top level
2021-03-05 11:07:50 -08:00
* Add support for building on Windows platforms
* Add support for general in_axes and out_axes in `jax.pmap`
* Add complex support for `jax.numpy.linalg.slogdet`
* Bug fixes:
* Fix higher-than-second order derivatives of `jax.numpy.sinc` at zero
* Fix some hard-to-hit bugs around symbolic zeros in transpose rules
* Breaking changes:
* `jax.experimental.optix` has been deleted, in favor of the standalone
`optax` Python package.
* indexing of JAX arrays with non-tuple sequences now raises a `TypeError`. This type of indexing
has been deprecated in Numpy since v1.16, and in JAX since v0.2.4.
See {jax-issue}`#4564`.
## jax 0.2.6 (Nov 18 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.2.5...jax-v0.2.6).
move changelog to top level
2021-03-05 11:07:50 -08:00
* New Features:
* Add support for shape-polymorphic tracing for the jax.experimental.jax2tf converter.
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
See [README.md](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/README.md).
move changelog to top level
2021-03-05 11:07:50 -08:00
* Breaking change cleanup
* Raise an error on non-hashable static arguments for jax.jit and
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
xla_computation. See [cb48f42](https://github.com/jax-ml/jax/commit/cb48f42).
move changelog to top level
2021-03-05 11:07:50 -08:00
* Improve consistency of type promotion behavior ({jax-issue}`#4744`):
* Adding a complex Python scalar to a JAX floating point number respects the precision of
the JAX float. For example, `jnp.float32(1) + 1j` now returns `complex64`, where previously
it returned `complex128`.
* Results of type promotion with 3 or more terms involving uint64, a signed int, and a third type
are now independent of the order of arguments. For example:
`jnp.result_type(jnp.uint64, jnp.int64, jnp.float16)` and
`jnp.result_type(jnp.float16, jnp.uint64, jnp.int64)` both return `float16`, where previously
the first returned `float64` and the second returned `float16`.
* The contents of the (undocumented) `jax.lax_linalg` linear algebra module
are now exposed publicly as `jax.lax.linalg`.
* `jax.random.PRNGKey` now produces the same results in and out of JIT compilation
({jax-issue}`#4877`).
This required changing the result for a given seed in a few particular cases:
* With `jax_enable_x64=False`, negative seeds passed as Python integers now return a different result
outside JIT mode. For example, `jax.random.PRNGKey(-1)` previously returned
`[4294967295, 4294967295]`, and now returns `[0, 4294967295]`. This matches the behavior in JIT.
* Seeds outside the range representable by `int64` outside JIT now result in an `OverflowError`
rather than a `TypeError`. This matches the behavior in JIT.
To recover the keys returned previously for negative integers with `jax_enable_x64=False`
outside JIT, you can use:
```
key = random.PRNGKey(-1).at[0].set(0xFFFFFFFF)
```
* DeviceArray now raises `RuntimeError` instead of `ValueError` when trying
to access its value while it has been deleted.
## jaxlib 0.1.58 (January 12ish 2021)
* Fixed a bug that meant JAX sometimes return platform-specific types (e.g.,
`np.cint`) instead of standard types (e.g., `np.int32`). (#4903)
* Fixed a crash when constant-folding certain int16 operations. (#4971)
* Added an `is_leaf` predicate to {func}`pytree.flatten`.
## jaxlib 0.1.57 (November 12 2020)
* Fixed manylinux2010 compliance issues in GPU wheels.
* Switched the CPU FFT implementation from Eigen to PocketFFT.
* Fixed a bug where the hash of bfloat16 values was not correctly initialized
and could change (#4651).
* Add support for retaining ownership when passing arrays to DLPack (#4636).
* Fixed a bug for batched triangular solves with sizes greater than 128 but not
a multiple of 128.
* Fixed a bug when performing concurrent FFTs on multiple GPUs (#3518).
* Fixed a bug in profiler where tools are missing (#4427).
* Dropped support for CUDA 10.0.
## jax 0.2.5 (October 27 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.2.4...jax-v0.2.5).
move changelog to top level
2021-03-05 11:07:50 -08:00
* Improvements:
* Ensure that `check_jaxpr` does not perform FLOPS. See {jax-issue}`#4650`.
* Expanded the set of JAX primitives converted by jax2tf.
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
See [primitives_with_limited_support.md](https://github.com/jax-ml/jax/blob/main/jax/experimental/jax2tf/primitives_with_limited_support.md).
move changelog to top level
2021-03-05 11:07:50 -08:00
## jax 0.2.4 (October 19 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.2.3...jax-v0.2.4).
move changelog to top level
2021-03-05 11:07:50 -08:00
* Improvements:
* Add support for `remat` to jax.experimental.host_callback. See {jax-issue}`#4608`.
* Deprecations
* Indexing with non-tuple sequences is now deprecated, following a similar deprecation in Numpy.
In a future release, this will result in a TypeError. See {jax-issue}`#4564`.
## jaxlib 0.1.56 (October 14, 2020)
## jax 0.2.3 (October 14 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.2.2...jax-v0.2.3).
move changelog to top level
2021-03-05 11:07:50 -08:00
* The reason for another release so soon is we need to temporarily roll back a
new jit fastpath while we look into a performance degradation
## jax 0.2.2 (October 13 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.2.1...jax-v0.2.2).
move changelog to top level
2021-03-05 11:07:50 -08:00
## jax 0.2.1 (October 6 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.2.0...jax-v0.2.1).
move changelog to top level
2021-03-05 11:07:50 -08:00
* Improvements:
* As a benefit of omnistaging, the host_callback functions are executed (in program
order) even if the result of the {py:func}`jax.experimental.host_callback.id_print`/
{py:func}`jax.experimental.host_callback.id_tap` is not used in the computation.
## jax (0.2.0) (September 23 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.77...jax-v0.2.0).
move changelog to top level
2021-03-05 11:07:50 -08:00
* Improvements:
pre-release omnistaging cleanup
2021-03-25 16:44:58 -07:00
* Omnistaging on by default. See {jax-issue}`#3370` and
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
[omnistaging](https://github.com/jax-ml/jax/blob/main/docs/design_notes/omnistaging.md)
move changelog to top level
2021-03-05 11:07:50 -08:00
## jax (0.1.77) (September 15 2020)
* Breaking changes:
* New simplified interface for {py:func}`jax.experimental.host_callback.id_tap` (#4101)
## jaxlib 0.1.55 (September 8, 2020)
* Update XLA:
* Fix bug in DLPackManagedTensorToBuffer (#4196)
## jax 0.1.76 (September 8, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.75...jax-v0.1.76).
move changelog to top level
2021-03-05 11:07:50 -08:00
## jax 0.1.75 (July 30, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.74...jax-v0.1.75).
move changelog to top level
2021-03-05 11:07:50 -08:00
* Bug Fixes:
* make jnp.abs() work for unsigned inputs (#3914)
* Improvements:
* "Omnistaging" behavior added behind a flag, disabled by default (#3370)
## jax 0.1.74 (July 29, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.73...jax-v0.1.74).
move changelog to top level
2021-03-05 11:07:50 -08:00
* New Features:
* BFGS (#3101)
DOC: many small fixes
2021-08-02 17:57:09 -07:00
* TPU support for half-precision arithmetic (#3878)
move changelog to top level
2021-03-05 11:07:50 -08:00
* Bug Fixes:
* Prevent some accidental dtype warnings (#3874)
* Fix a multi-threading bug in custom derivatives (#3845, #3869)
* Improvements:
* Faster searchsorted implementation (#3873)
* Better test coverage for jax.numpy sorting algorithms (#3836)
## jaxlib 0.1.52 (July 22, 2020)
* Update XLA.
## jax 0.1.73 (July 22, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.72...jax-v0.1.73).
move changelog to top level
2021-03-05 11:07:50 -08:00
* The minimum jaxlib version is now 0.1.51.
* New Features:
* jax.image.resize. (#3703)
* hfft and ihfft (#3664)
* jax.numpy.intersect1d (#3726)
* jax.numpy.lexsort (#3812)
* `lax.scan` and the `scan` primitive support an `unroll`
parameter for loop unrolling when lowering to XLA
({jax-issue}`#3738`).
* Bug Fixes:
* Fix reduction repeated axis error (#3618)
* Fix shape rule for lax.pad for input dimensions of size 0. (#3608)
* make psum transpose handle zero cotangents (#3653)
* Fix shape error when taking JVP of reduce-prod over size 0 axis. (#3729)
* Support differentiation through jax.lax.all_to_all (#3733)
* address nan issue in jax.scipy.special.zeta (#3777)
* Improvements:
* Many improvements to jax2tf
* Reimplement argmin/argmax using a single pass variadic reduction. (#3611)
* Enable XLA SPMD partitioning by default. (#3151)
* Add support for 0d transpose convolution (#3643)
* Make LU gradient work for low-rank matrices (#3610)
* support multiple_results and custom JVPs in jet (#3657)
* Generalize reduce-window padding to support (lo, hi) pairs. (#3728)
* Implement complex convolutions on CPU and GPU. (#3735)
* Make jnp.take work for empty slices of empty arrays. (#3751)
* Relax dimension ordering rules for dot_general. (#3778)
* Enable buffer donation for GPU. (#3800)
* Add support for base dilation and window dilation to reduce window op… (#3803)
## jaxlib 0.1.51 (July 2, 2020)
* Update XLA.
* Add new runtime support for host_callback.
## jax 0.1.72 (June 28, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.71...jax-v0.1.72).
move changelog to top level
2021-03-05 11:07:50 -08:00
* Bug fixes:
* Fix an odeint bug introduced in the previous release, see
{jax-issue}`#3587`.
## jax 0.1.71 (June 25, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.70...jax-v0.1.71).
move changelog to top level
2021-03-05 11:07:50 -08:00
* The minimum jaxlib version is now 0.1.48.
* Bug fixes:
* Allow `jax.experimental.ode.odeint` dynamics functions to close over
values with respect to which we're differentiating
{jax-issue}`#3562`.
## jaxlib 0.1.50 (June 25, 2020)
* Add support for CUDA 11.0.
* Drop support for CUDA 9.2 (we only maintain support for the last four CUDA
versions.)
* Update XLA.
## jaxlib 0.1.49 (June 19, 2020)
* Bug fixes:
* Fix build issue that could result in slow compiles
(<https://github.com/tensorflow/tensorflow/commit/f805153a25b00d12072bd728e91bb1621bfcf1b1>)
## jaxlib 0.1.48 (June 12, 2020)
* New features:
* Adds support for fast traceback collection.
* Adds preliminary support for on-device heap profiling.
* Implements `np.nextafter` for `bfloat16` types.
* Complex128 support for FFTs on CPU and GPU.
Fix Typos in CHANGELOG.md This PR fixes the typos in Change log documentation
2024-03-12 13:57:07 +05:30
* Bug fixes:
move changelog to top level
2021-03-05 11:07:50 -08:00
* Improved float64 `tanh` accuracy on GPU.
* float64 scatters on GPU are much faster.
* Complex matrix multiplication on CPU should be much faster.
* Stable sorts on CPU should actually be stable now.
* Concurrency bug fix in CPU backend.
## jax 0.1.70 (June 8, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.69...jax-v0.1.70).
move changelog to top level
2021-03-05 11:07:50 -08:00
* New features:
* `lax.switch` introduces indexed conditionals with multiple
branches, together with a generalization of the `cond`
primitive
{jax-issue}`#3318`.
## jax 0.1.69 (June 3, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.68...jax-v0.1.69).
move changelog to top level
2021-03-05 11:07:50 -08:00
## jax 0.1.68 (May 21, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.67...jax-v0.1.68).
move changelog to top level
2021-03-05 11:07:50 -08:00
* New features:
* {func}`lax.cond` supports a single-operand form, taken as the argument
to both branches
{jax-issue}`#2993`.
* Notable changes:
* The format of the `transforms` keyword for the {func}`jax.experimental.host_callback.id_tap`
primitive has changed {jax-issue}`#3132`.
## jax 0.1.67 (May 12, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.66...jax-v0.1.67).
move changelog to top level
2021-03-05 11:07:50 -08:00
* New features:
* Support for reduction over subsets of a pmapped axis using `axis_index_groups`
{jax-issue}`#2382`.
* Experimental support for printing and calling host-side Python function from
compiled code. See [id_print and id_tap](https://jax.readthedocs.io/en/latest/jax.experimental.host_callback.html)
({jax-issue}`#3006`).
* Notable changes:
* The visibility of names exported from {mod}`jax.numpy` has been
tightened. This may break code that was making use of names that were
previously exported accidentally.
## jaxlib 0.1.47 (May 8, 2020)
* Fixes crash for outfeed.
## jax 0.1.66 (May 5, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.65...jax-v0.1.66).
move changelog to top level
2021-03-05 11:07:50 -08:00
* New features:
* Support for `in_axes=None` on {func}`pmap`
{jax-issue}`#2896`.
## jaxlib 0.1.46 (May 5, 2020)
* Fixes crash for linear algebra functions on Mac OS X (#432).
* Fixes an illegal instruction crash caused by using AVX512 instructions when
an operating system or hypervisor disabled them (#2906).
## jax 0.1.65 (April 30, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.64...jax-v0.1.65).
move changelog to top level
2021-03-05 11:07:50 -08:00
* New features:
* Differentiation of determinants of singular matrices
{jax-issue}`#2809`.
* Bug fixes:
* Fix {func}`odeint` differentiation with respect to time of ODEs with
time-dependent dynamics {jax-issue}`#2817`,
also add ODE CI testing.
* Fix {func}`lax_linalg.qr` differentiation
{jax-issue}`#2867`.
## jaxlib 0.1.45 (April 21, 2020)
* Fixes segfault: {jax-issue}`#2755`
* Plumb is_stable option on Sort HLO through to Python.
## jax 0.1.64 (April 21, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.63...jax-v0.1.64).
move changelog to top level
2021-03-05 11:07:50 -08:00
* New features:
* Add syntactic sugar for functional indexed updates
{jax-issue}`#2684`.
* Add {func}`jax.numpy.linalg.multi_dot` {jax-issue}`#2726`.
* Add {func}`jax.numpy.unique` {jax-issue}`#2760`.
* Add {func}`jax.numpy.rint` {jax-issue}`#2724`.
* Add {func}`jax.numpy.rint` {jax-issue}`#2724`.
* Add more primitive rules for {func}`jax.experimental.jet`.
* Bug fixes:
* Fix {func}`logaddexp` and {func}`logaddexp2` differentiation at zero {jax-issue}`#2107`.
* Improve memory usage in reverse-mode autodiff without {func}`jit`
{jax-issue}`#2719`.
* Better errors:
* Improves error message for reverse-mode differentiation of {func}`lax.while_loop`
{jax-issue}`#2129`.
## jaxlib 0.1.44 (April 16, 2020)
* Fixes a bug where if multiple GPUs of different models were present, JAX
would only compile programs suitable for the first GPU.
* Bugfix for `batch_group_count` convolutions.
* Added precompiled SASS for more GPU versions to avoid startup PTX compilation
hang.
## jax 0.1.63 (April 12, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.62...jax-v0.1.63).
move changelog to top level
2021-03-05 11:07:50 -08:00
* Added `jax.custom_jvp` and `jax.custom_vjp` from {jax-issue}`#2026`, see the [tutorial notebook](https://jax.readthedocs.io/en/latest/notebooks/Custom_derivative_rules_for_Python_code.html). Deprecated `jax.custom_transforms` and removed it from the docs (though it still works).
* Add `scipy.sparse.linalg.cg` {jax-issue}`#2566`.
* Changed how Tracers are printed to show more useful information for debugging {jax-issue}`#2591`.
* Made `jax.numpy.isclose` handle `nan` and `inf` correctly {jax-issue}`#2501`.
* Added several new rules for `jax.experimental.jet` {jax-issue}`#2537`.
* Fixed `jax.experimental.stax.BatchNorm` when `scale`/`center` isn't provided.
* Fix some missing cases of broadcasting in `jax.numpy.einsum` {jax-issue}`#2512`.
Finalize deprecations of several APIs PiperOrigin-RevId: 633634215
2024-05-14 10:39:43 -07:00
* Implement `jax.numpy.cumsum` and `jax.numpy.cumprod` in terms of a parallel prefix scan {jax-issue}`#2596` and make `reduce_prod` differentiable to arbitrary order {jax-issue}`#2597`.
move changelog to top level
2021-03-05 11:07:50 -08:00
* Add `batch_group_count` to `conv_general_dilated` {jax-issue}`#2635`.
* Add docstring for `test_util.check_grads` {jax-issue}`#2656`.
* Add `callback_transform` {jax-issue}`#2665`.
* Implement `rollaxis`, `convolve`/`correlate` 1d & 2d, `copysign`,
`trunc`, `roots`, and `quantile`/`percentile` interpolation options.
## jaxlib 0.1.43 (March 31, 2020)
* Fixed a performance regression for Resnet-50 on GPU.
## jax 0.1.62 (March 21, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.61...jax-v0.1.62).
move changelog to top level
2021-03-05 11:07:50 -08:00
* JAX has dropped support for Python 3.5. Please upgrade to Python 3.6 or newer.
* Removed the internal function `lax._safe_mul`, which implemented the
convention `0. * nan == 0.`. This change means some programs when
differentiated will produce nans when they previously produced correct
values, though it ensures nans rather than silently incorrect results are
produced for other programs. See #2447 and #1052 for details.
* Added an `all_gather` parallel convenience function.
* More type annotations in core code.
## jaxlib 0.1.42 (March 19, 2020)
* jaxlib 0.1.41 broke cloud TPU support due to an API incompatibility. This
release fixes it again.
* JAX has dropped support for Python 3.5. Please upgrade to Python 3.6 or newer.
## jax 0.1.61 (March 17, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.60...jax-v0.1.61).
move changelog to top level
2021-03-05 11:07:50 -08:00
* Fixes Python 3.5 support. This will be the last JAX or jaxlib release that
supports Python 3.5.
## jax 0.1.60 (March 17, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.59...jax-v0.1.60).
move changelog to top level
2021-03-05 11:07:50 -08:00
* New features:
* {py:func}`jax.pmap` has `static_broadcast_argnums` argument which allows
the user to specify arguments that should be treated as compile-time
constants and should be broadcasted to all devices. It works analogously to
`static_argnums` in {py:func}`jax.jit`.
* Improved error messages for when tracers are mistakenly saved in global state.
* Added {py:func}`jax.nn.one_hot` utility function.
* Added {mod}`jax.experimental.jet` for exponentially faster
higher-order automatic differentiation.
* Added more correctness checking to arguments of {py:func}`jax.lax.broadcast_in_dim`.
* The minimum jaxlib version is now 0.1.41.
## jaxlib 0.1.40 (March 4, 2020)
* Adds experimental support in Jaxlib for TensorFlow profiler, which allows
tracing of CPU and GPU computations from TensorBoard.
* Includes prototype support for multihost GPU computations that communicate via
NCCL.
* Improves performance of NCCL collectives on GPU.
* Adds TopK, CustomCallWithoutLayout, CustomCallWithLayout, IGammaGradA and
RandomGamma implementations.
* Supports device assignments known at XLA compilation time.
## jax 0.1.59 (February 11, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/jax-v0.1.58...jax-v0.1.59).
move changelog to top level
2021-03-05 11:07:50 -08:00
* Breaking changes
* The minimum jaxlib version is now 0.1.38.
* Simplified {py:class}`Jaxpr` by removing the `Jaxpr.freevars` and
`Jaxpr.bound_subjaxprs`. The call primitives (`xla_call`, `xla_pmap`,
`sharded_call`, and `remat_call`) get a new parameter `call_jaxpr` with a
fully-closed (no `constvars`) jaxpr. Also, added a new field `call_primitive`
to primitives.
* New features:
* Reverse-mode automatic differentiation (e.g. `grad`) of `lax.cond`, making it
now differentiable in both modes ({jax-issue}`#2091`)
* JAX now supports DLPack, which allows sharing CPU and GPU arrays in a
zero-copy way with other libraries, such as PyTorch.
* JAX GPU DeviceArrays now support `__cuda_array_interface__`, which is another
zero-copy protocol for sharing GPU arrays with other libraries such as CuPy
and Numba.
* JAX CPU device buffers now implement the Python buffer protocol, which allows
zero-copy buffer sharing between JAX and NumPy.
* Added JAX_SKIP_SLOW_TESTS environment variable to skip tests known as slow.
## jaxlib 0.1.39 (February 11, 2020)
* Updates XLA.
## jaxlib 0.1.38 (January 29, 2020)
* CUDA 9.0 is no longer supported.
* CUDA 10.2 wheels are now built by default.
## jax 0.1.58 (January 28, 2020)
Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax PiperOrigin-RevId: 676843138
2024-09-20 07:51:48 -07:00
* [GitHub commits](https://github.com/jax-ml/jax/compare/46014da21...jax-v0.1.58).
move changelog to top level
2021-03-05 11:07:50 -08:00
* Breaking changes
* JAX has dropped Python 2 support, because Python 2 reached its end of life on
January 1, 2020. Please update to Python 3.5 or newer.
* New features
> > * Forward-mode automatic differentiation (`jvp`) of while loop
> ({jax-issue}`#1980`)
> * New NumPy and SciPy functions:
Update jaxlib version etc.
2021-03-09 17:55:40 -08:00
>
move changelog to top level
2021-03-05 11:07:50 -08:00
> * {py:func}`jax.numpy.fft.fft2`
> * {py:func}`jax.numpy.fft.ifft2`
> * {py:func}`jax.numpy.fft.rfft`
> * {py:func}`jax.numpy.fft.irfft`
> * {py:func}`jax.numpy.fft.rfft2`
> * {py:func}`jax.numpy.fft.irfft2`
> * {py:func}`jax.numpy.fft.rfftn`
> * {py:func}`jax.numpy.fft.irfftn`
> * {py:func}`jax.numpy.fft.fftfreq`
> * {py:func}`jax.numpy.fft.rfftfreq`
> * {py:func}`jax.numpy.linalg.matrix_rank`
> * {py:func}`jax.numpy.linalg.matrix_power`
> * {py:func}`jax.scipy.special.betainc`
> * Batched Cholesky decomposition on GPU now uses a more efficient batched
> kernel.
### Notable bug fixes
* With the Python 3 upgrade, JAX no longer depends on `fastcache`, which should
help with installation.
Reference in New Issue Copy Permalink