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.
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.
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](https://github.com/google/jax/issues/6501))
{func}jax.numpy.unique now supports the axis argument ({jax-issue}[#6532](https://github.com/google/jax/issues/6532)).
{func}jax.experimental.host_callback.call now supports pjit.pjit ({jax-issue}[#6569](https://github.com/google/jax/issues/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
... (truncated)
Commits
bea1807 Merge pull request #6767 from skye:workspace
Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting @dependabot rebase.
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR:
- `@dependabot rebase` will rebase this PR
- `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it
- `@dependabot merge` will merge this PR after your CI passes on it
- `@dependabot squash and merge` will squash and merge this PR after your CI passes on it
- `@dependabot cancel merge` will cancel a previously requested merge and block automerging
- `@dependabot reopen` will reopen this PR if it is closed
- `@dependabot close` will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually
- `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
- `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
- `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
Bumps jaxlib from 0.1.61 to 0.1.67.
Changelog
Sourced from jaxlib's changelog.
... (truncated)
Commits
bea1807
Merge pull request #6767 from skye:workspacec5e199a
Update WORKSPACE for jaxlib 0.1.67149f7fa
Tighten the numerical tolerances for jax2tf.afe5ec3
Improve accuracy of the jax2tf convolution conversion.e049d89
Merge pull request #6757 from LenaMartens:changelist/3737865346a28351
Merge pull request #6760 from gnecula:tf_limitationsa08cdb3
[jax2tf] Update the limitations for unsupported primitives25cc3ec
Merge pull request #6739 from jakevdp:sparse-op-jvp2bcae23
fix flakes & test tolerances25ff2f4
propagate symbolic zeros in sparse op JVPsDependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting
@dependabot rebase
.Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot merge` will merge this PR after your CI passes on it - `@dependabot squash and merge` will squash and merge this PR after your CI passes on it - `@dependabot cancel merge` will cancel a previously requested merge and block automerging - `@dependabot reopen` will reopen this PR if it is closed - `@dependabot close` will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually - `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself) - `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)