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).
jax.random.choice and jax.random.permutation now support
multidimensional arrays and an optional axis argument ({jax-issue}[#8158](https://github.com/google/jax/issues/8158))
Breaking changes:
jax.numpy.take and jax.numpy.take_along_axis now require array-like inputs
(see {jax-issue}[#7737](https://github.com/google/jax/issues/7737))
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:
pip install --upgrade pip
Installs the wheel compatible with CUDA 11 and cuDNN 8.2 or newer.
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
... (truncated)
Commits
a7c9b6d Update jax version number for jax release.
185d7a9 Delete xla_bridge.xla.dtype_to_etype, replace it with jax.interpreters.xla.dt...
ee752b3 Use cuda11_cudnn82 instead of cuda=11,cudnn=82 because the latter one is ...
a91cb81 Merge pull request #8269 from jakevdp:sparse-dtype
4d8bce1 Add a default cuda installation path and more explicit installation paths for...
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 jax from 0.2.8 to 0.2.24.
Release notes
Sourced from jax's releases.
Changelog
Sourced from jax's changelog.
... (truncated)
Commits
a7c9b6d
Update jax version number for jax release.185d7a9
Delete xla_bridge.xla.dtype_to_etype, replace it with jax.interpreters.xla.dt...ee752b3
Usecuda11_cudnn82
instead ofcuda=11,cudnn=82
because the latter one is ...a91cb81
Merge pull request #8269 from jakevdp:sparse-dtype4d8bce1
Add a default cuda installation path and more explicit installation paths for...6bd0c72
Silence some pytype errors.fa89f34
[JAX] Port lax translation rules to updated XLA translation rule API.fa5d4a2
Merge pull request #8268 from jakevdp:sparsify-weak-typesfbdeff0
Update the workspace file48bbdbc
Change jax.core.DropVar to be a non-singleton.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)