jax

• The release was yanked.

• Changes
  • {func}jax.numpy.linalg.pinv now supports the hermitian option.
  • {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 {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.
• Breaking Changes
  • Deleted the jax_experimental_name_stack config option.
  • 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.

• Changes
  • JAX should be faster to import. We now import scipy lazily, which accounted for a significant fraction of JAX's import time.
  • 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.
    • Added {func}jax.scipy.stats.mode.
  • 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
  • {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)
  • 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.
• Deprecations
  • jax.sharding.MeshPspecSharding has been renamed to jax.sharding.NamedSharding. jax.sharding.MeshPspecSharding name will be removed in 3 months.

• Changes
  • Update Colab TPU driver version for new jaxlib release.

• Changes
  • 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.
• Deprecations
  • Several test utilities deprecated in JAX v0.3.8 are now removed from {mod}jax.test_util.

• 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.

Notable changes:

• Adds missing .pyi files that were missing from the previous release (#12536).
• Fixes an incompatibility between jax 0.3.19 and the libtpu version it pinned (#12550). Requires jaxlib 0.3.20.
• Fix incorrect pip url in setup.py comment (#12528).

Fixes the required jaxlib version

• GitHub commits.
• Changes
  • Ahead-of-time lowering and compilation functionality (tracked in {jax-issue}#7733) is stable and public. See the overview and the API docs for {mod}jax.stages.
  • 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.
• Breaking changes
  • jax._src is no longer imported into the from the public jax namespace. This may break users that were using JAX internals.
  • 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.

• GitHub commits.
• Bugs
  • Fix corner case issue in gradient of lax.pow with an exponent of zero (#12041)
• Breaking changes
  • jax.checkpoint, also known as jax.remat, no longer supports the concrete option, following the previous version's deprecation; see JEP 11830.
• Changes
  • Added jax.pure_callback that enables calling back to pure Python functions from compiled functions (e.g. functions decorated with jax.jit or jax.pmap).
• Deprecations:
  • The deprecated DeviceArray.tile() method has been removed. Use jax.numpy.tile (#11944).
  • DeviceArray.to_py() has been deprecated. Use np.asarray(x) instead.

• GitHub commits.
• Breaking changes
  • Support for NumPy 1.19 has been dropped, per the deprecation policy. Please upgrade to NumPy 1.20 or newer.
• Changes
  • Added jax.debug that includes utilities for runtime value debugging such at jax.debug.print and jax.debug.breakpoint.
  • Added new documentation for runtime value debugging
• Deprecations
  • jax.mask jax.shapecheck APIs have been removed. See #11557.
  • jax.experimental.loops has been removed. See #10278 for an alternative API.
  • jax.tree_util.tree_multimap has been removed. It has been deprecated since JAX release 0.3.5, and jax.tree_util.tree_map is a direct replacement.
  • Removed jax.experimental.stax; it has long been a deprecated alias of jax.example_libraries.stax.
  • Removed jax.experimental.optimizers; it has long been a deprecated alias of jax.example_libraries.optimizers.
  • jax.checkpoint, also known as jax.remat, has a new implementation switched on by default, meaning the old implementation is deprecated; see JEP 11830.

• GitHub commits.
• Changes
  • JaxTestCase and JaxTestLoader have been removed from jax.test_util. These classes have been deprecated since v0.3.1 ({jax-issue}#11248).
  • Added {class}jax.scipy.gaussian_kde ({jax-issue}#11237).
  • 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).
  • 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:
    • {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
  • 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.

• GitHub commits.
• Breaking changes
  • {func}jax.experimental.compilation_cache.initialize_cache does not support max_cache_size_ bytes anymore and will not get that as an input.
  • JAX_PLATFORMS now raises an exception when platform initialization fails.
• Changes
  • Fixed compatibility problems with NumPy 1.23.
  • {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".
  • 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().
    • Calling pickle on a traced array now results in an explicit ConcretizationTypeError.
  • 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).
  • {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 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.
  • 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 {func}jax.random.generalized_normal.
  • Added {func}jax.random.ball.
  • Added {func}jax.default_device.
  • Added a python -m jax.collect_profile script to manually capture program traces as an alternative to the TensorBoard UI.
  • Added a jax.named_scope context manager that adds profiler metadata to Python programs (similar to jax.named_call).
  • 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.
  • {func}jax.experimental.compilation_cache.initialize_cache now supports gcs bucket path as input.
  • Added {func}jax.scipy.stats.gennorm.
  • {func}jax.numpy.roots is now better behaved when strip_zeros=False when coefficients have leading zeros ({jax-issue}#11215).

• GitHub commits.

• Changes
  • Fixes https://github.com/google/jax/pull/10717

• Changes
  • {func}jax.lax.eigh now accepts an optional sort_eigenvalues argument that allows users to opt out of eigenvalue sorting on TPU.
• Deprecations
  • 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.
  • {func}jax.scipy.linalg.polar_unitary, which was a JAX extension to the scipy API, is deprecated. Use {func}jax.scipy.linalg.polar instead.

• GitHub commits.

• Changes
  • Added support for fully asynchronous checkpointing for GlobalDeviceArray.

• GitHub commits.
• Changes
  • {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.
  • {func}jax.scipy.cluster.vq.vq has been added.
  • 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.
  • {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)
  • {func}jax.numpy.take_along_axis now takes an optional mode parameter 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".
  • {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. 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.
  • {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.
  • {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 {func}jax.random.orthogonal.
• 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. Most of these utilites can be replaced by calls to standard python & numpy testing utilities found 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.

• Fixed a performance problem if the indices passed to jax.numpy.take_along_axis were broadcasted (#10281).
• jax.scipy.special.expit and jax.scipy.special.logit now require their arguments to be scalars or JAX arrays. They also now promote integer arguments to floating point.
• The DeviceArray.tile() method is deprecated, because numpy arrays do not have a tile() method. As a replacement for this, use jax.numpy.tile (#10266).

• Changes:
  • Upgraded libtpu wheel to the fixed version. Fixes #10218.

Changes

• added jax.random.loggamma & improved behavior of jax.random.beta and jax.random.dirichlet for small parameter values (#9906).
• the private lax_numpy submodule is no longer exposed in the jax.numpy namespace (#10029).
• added array creation routines jax.numpy.frombuffer, jax.numpy.fromfunction, and jax.numpy.fromstring (#10049).
• DeviceArray.copy() now returns a DeviceArray rather than a np.ndarray (#10069)
• added jax.scipy.linalg.rsf2csf
• Deprecations:
  • jax.nn.normalize is being deprecated. Use jax.nn.standardize instead (#9899).
  • jax.tree_util.tree_multimap is deprecated. Use jax.tree_util.tree_map instead (#5746).
  • jax.experimental.sharded_jit is deprecated. Use pjit instead.

Fix a bug introduced in #9923.

• GitHub commits.

• GitHub commits.
• 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 instead, e.g., x.at[idx].set(y).
  • Moved jax.experimental.ann.approx_*_k into jax.lax. These functions are optimized alternatives to jax.lax.top_k.
  • {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).
  • The standard jax[tpu] install can now be used with Cloud TPU v4 VMs.
  • pjit now works on CPU (in addition to previous TPU and GPU support).

• 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 numpy.testing.assert_allclose(), which work directly with JAX arrays (#9620 ).
  • jax.test_util.JaxTestCase now sets jax_numpy_rank_promotion='raise' by default (#9562 ). To recover the previous behavior, use the new jax.test_util.with_config decorator:

    @jtu.with_config(jax_numpy_rank_promotion='allow')
    class MyTestCase(jtu.JaxTestCase):
      ...
    

  • Added jax.scipy.linalg.schur, jax.scipy.linalg.sqrtm, jax.scipy.signal.csd, jax.scipy.signal.stft, jax.scipy.signal.welch.

• Changes
  • jax version has been bumped to 0.3.0. Please see the design doc for the explanation.

• GitHub commits.
  • jax.jit(f).lower(...).compiler_ir() now defaults to the MHLO dialect if no dialect= is passed.
  • The jax.jit(f).lower(...).compiler_ir(dialect='mhlo') now returns an MLIR ir.Module object instead of its string representation.

• GitHub commits.

• 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.
  • The host_callback primitives have been simplified to drop the 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 environment variable, or the --flax_host_callback_ad_transforms flag. Additionally, added documentation for how to implement the old behavior using JAX custom AD APIs ({jax-issue}#8678).
  • 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).
  • {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).

• Bug fixes:

  • host_callback now supports ad_checkpoint.checkpoint ({jax-issue}#8907).

• New features:

  • add jax.block_until_ready ({jax-issue}`#8941)
  • 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.
  • Added jax.ensure_compile_time_eval to the public api ({jax-issue}#7987).
  • 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).

• Bug fixes:

• Out-of-bounds indices to jax.ops.segment_sum will now be handled with FILL_OR_DROP semantics, as documented. This primarily afects the reverse-mode derivative, where gradients corresponding to out-of-bounds indices will now be returned as 0. (#8634).

• jax2tf will force the converted code to use XLA for the code fragments under jax.jit, e.g., most jax.numpy functions (#7839).

• New features:

  • (Experimental) jax.distributed.initialize exposes multi-host GPU backend.
  • jax.random.permutation supports new independent keyword argument ({jax-issue}#8430)

• Breaking changes

  • Moved jax.experimental.stax to jax.example_libraries.stax
  • Moved jax.experimental.optimizers to jax.example_libraries.optimizers

• New features:

  • Added jax.lax.linalg.qdwh.

• New features:
  • jax.random.choice and jax.random.permutation now support multidimensional arrays and an optional axis argument (#8158)
• Breaking changes:
  • jax.numpy.take and jax.numpy.take_along_axis now require array-like inputs (see #7737)

• GitHub commits.
• 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.

• 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 instead, e.g., x.at[idx].set(y). For now, these functions produce a DeprecationWarning.
• New features:
  • 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).
  • jax.numpy.unique now supports an optional fill_value argument ({jax-issue}#8121)

• New features:

  • Added jax.numpy.insert implementation (#7936 ).

• Breaking Changes

  • 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.
  • jax.partial, jax.lax.partial, and jax.util.partial were accidental exports that have now been removed. Use functools.partial from the Python standard library instead.
  • Boolean scalar indices now raise a TypeError; previously this silently returned wrong results (#7925 ).
  • Many more jax.numpy functions now require array-like inputs, and will error if passed a list (#7747 #7802 #7907 ). See #7737 for a discussion of the rationale behind this change.
  • 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.
  • 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 (#7927).

• GitHub commits.
• 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)

• GitHub commits.
• Breaking changes:

  • Support for NumPy 1.17 has been dropped, per the deprecation policy. Please upgrade to a supported NumPy version.

  • 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)).

• 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)
• Bug fixes:
  • Some leaked trace errors from the previous release ({jax-issue}#7613)

• GitHub commits.

• Breaking changes:

  • Support for Python 3.6 has been dropped, per the deprecation policy. Please upgrade to a supported Python version.
  • The minimum jaxlib version is now 0.1.69.
  • The backend argument to {py:func}jax.dlpack.from_dlpack has been removed.

• New features:

  • Added a polar decomposition ({py:func}jax.scipy.linalg.polar).

• Bug fixes:

  • Tightened the checks for lax.argmin and lax.argmax to ensure they are not used with an invalid axis value, or with an empty reduction dimension. ({jax-issue}#7196)

• GitHub commits.
• 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.
• New features:
  • New SciPy function {py:func}jax.scipy.special.sph_harm.
  • 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).

• GitHub commits.

• GitHub commits.

• New features:

  • #7042 Turned on TFRT CPU backend with significant dispatch performance improvements on CPU.
  • The {func}jax2tf.convert supports inequalities and min/max for booleans ({jax-issue}#6956).
  • New SciPy function {py:func}jax.scipy.special.lpmn_values.

• Breaking changes:

  • Support for NumPy 1.16 has been dropped, per the deprecation policy.

• Bug fixes:

  • Fixed bug that prevented round-tripping from JAX to TF and back: jax2tf.call_tf(jax2tf.convert) ({jax-issue}#6947).

• GitHub commits.

• New features:

  • The {func}jax2tf.convert now has support for pjit and sharded_jit.
  • 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.
  • The {func}jax2tf.convert supports shape polymorphism even when the unknown dimensions are used in arithmetic operations, e.g., jnp.reshape(-1) ({jax-issue}#6827).
  • 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.
  • New SciPy function {py:func}jax.scipy.special.lpmn.

• Bug fixes:

  • 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).
  • The {func}jax2tf.convert now scopes the enable_xla conversion parameter properly to apply only during the just-in-time conversion ({jax-issue}#6720).
  • 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).
  • The {func}jax2tf.convert now has support for inequality comparisons and min/max for complex numbers ({jax-issue}#6892).

• GitHub commits.

• 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.
  • {func}jax.nonzero has a new optional size argument that allows it to be used within jit ({jax-issue}#6501)
  • {func}jax.numpy.unique now supports the axis argument ({jax-issue}#6532).
  • {func}jax.experimental.host_callback.call now supports pjit.pjit ({jax-issue}#6569).
  • Added {func}jax.scipy.linalg.eigh_tridiagonal that computes the eigenvalues of a tridiagonal matrix. Only eigenvalues are supported at present.
  • 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.

• Breaking changes:

  • 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.
  • 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.

• Bug fixes:

  • The {func}jax2tf.convert now works in presence of gradients for functions with integer inputs ({jax-issue}#6360).
  • Fixed assertion failure in {func}jax2tf.call_tf when used with captured tf.Variable ({jax-issue}#6572).

• GitHub commits.
• New features
  • New profiling APIs: {func}jax.profiler.start_trace, {func}jax.profiler.stop_trace, and {func}jax.profiler.trace
  • {func}jax.lax.reduce is now differentiable.
• Breaking changes:
  • The minimum jaxlib version is now 0.1.64.
  • 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
  • Omnistaging can no longer be disabled. See omnistaging for more information.
  • 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).
  • Outside X64 mode, Python integers outside the range representable by int32 will now lead to an OverflowError rather than having their value silently truncated.
• Bug fixes:
  • host_callback now supports empty arrays in arguments and results ({jax-issue}#6262).
  • {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)

• GitHub commits.

• New features:

  • #6112 added context managers: jax.enable_checks, jax.check_tracer_leaks, jax.debug_nans, jax.debug_infs, jax.log_compiles.
  • #6085 added jnp.delete

• Bug fixes:

  • #6136 generalized jax.flatten_util.ravel_pytree to handle integer dtypes.
  • #6129 fixed a bug with handling some constants like enum.IntEnums
  • #6145 fixed batching issues with incomplete beta functions
  • #6014 fixed H2D transfers during tracing
  • #6165 avoids OverflowErrors when converting some large Python integers to floats

• Breaking changes:

  • The minimum jaxlib version is now 0.1.62.

• GitHub commits.
• 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) and README).
  • 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, :funcjax.numpy.ravel, {func}jax.numpy.repeat, {func}jax.numpy.reshape. In general, {mod}jax.numpy functions should be used with scalars or array arguments.

• GitHub commits.
• 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.

• GitHub commits.
• 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)
  • host_callback.id_tap now works for jax.pmap also. There is an 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:
  • New flag for debugging inf, analogous to that for NaN ({jax-issue}#5224).

• GitHub commits.
• New features:
  • Add jax.device_put_replicated
  • Add multi-host support to jax.experimental.sharded_jit
  • Add support for differentiating eigenvalues computed by jax.numpy.linalg.eig
  • 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.

• GitHub commits.

• New Features:

  • Add support for shape-polymorphic tracing for the jax.experimental.jax2tf converter. See README.md.

• Breaking change cleanup

  • Raise an error on non-hashable static arguments for jax.jit and xla_computation. See cb48f42.

  • 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.

• GitHub commits.
• Improvements:
  • Ensure that check_jaxpr does not perform FLOPS. See {jax-issue}#4650.
  • Expanded the set of JAX primitives converted by jax2tf. See primitives_with_limited_support.md.

• GitHub commits.

• 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.

• GitHub commits.
• 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

• GitHub commits.

• GitHub commits.
• 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.

• GitHub commits.
• Improvements:
  • Omnistaging on by default. See {jax-issue}#3370 and omnistaging

• Breaking changes:
  • New simplified interface for {py:func}jax.experimental.host_callback.id_tap (#4101)

• 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.
  • With jaxlib 0.1.76, JAX uses the MHLO MLIR dialect as its primary target compiler IR by default.
• Breaking changes
  • Support for NumPy 1.18 has been dropped, per the deprecation policy. Please upgrade to a supported NumPy version.
• 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).