jax

• Support for CUDA 10.2 and CUDA 10.1 has been dropped. Jaxlib now supports CUDA 11.1+.
• Bug fixes:
  • Fixes https://github.com/jax-ml/jax/issues/7461, which caused wrong outputs on all platforms due to incorrect buffer aliasing inside the XLA compiler.

• Support for CUDA 11.0 and CUDA 10.1 has been dropped. Jaxlib now supports CUDA 10.2 and CUDA 11.1+.

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

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

• The host_callback mechanism now uses one thread per local device for 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.

• Fixed bug in TFRT CPU backend that gets nans when transfer TPU buffer to CPU.

• 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 ({jax-issue}#3006).

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

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

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

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

• Differentiation of determinants of singular matrices {jax-issue}#2809.

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

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

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

• Improves error message for reverse-mode differentiation of {func}lax.while_loop {jax-issue}#2129.

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

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

GitHub commits.

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.

• Fix bug in DLPackManagedTensorToBuffer (#4196)

• Fix build issue that could result in slow compiles (https://github.com/tensorflow/tensorflow/commit/f805153a25b00d12072bd728e91bb1621bfcf1b1)

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

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