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.
{func}jax.random.gamma has been re-factored to a more efficient algorithm
with more robust endpoint behavior ({jax-issue}[#16779](https://github.com/google/jax/issues/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,
{func}jax.random.beta, {func}jax.random.chisquare, {func}jax.random.dirichlet,
{func}jax.random.generalized_normal, {func}jax.random.loggamma, {func}jax.random.t).
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.
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)
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.
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.
You can trigger a rebase of this PR 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)
Note
Automatic rebases have been disabled on this pull request as it has been open for over 30 days.
Updates the requirements on jax[cpu] to permit the latest version.
Changelog
Sourced from jax[cpu]'s changelog.
... (truncated)
Commits
88a60b8
Merge pull request #16870 from skye:versione132a0e
Slightly downgrade xla version to avoid PJRT C API incompatc75e85d
Merge pull request #16869 from skye:versiona480aa8
Work around pytype error.0b24b2b
Update WORKSPACE and setup.py in preparation for jax/jaxlib 0.4.14 release7c0ef86
Merge pull request #16433 from JGameCreation:patch-176cda0a
Update flags to use the ABSL typed flag API.f35f226
Merge pull request #16865 from google:xla0cd0253
Update XLA version to fix build failure.a03d6e6
Move _tpu_ext.cc to jaxlib/mlir/_mlir_libs and set RPATH correctlyYou can trigger a rebase of this PR 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)