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)).
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](https://github.com/google/jax/issues/6360)).
Fixed assertion failure in {func}jax2tf.call_tf when used with captured
tf.Variable ({jax-issue}[#6572](https://github.com/google/jax/issues/6572)).
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.
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](https://github.com/google/jax/issues/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.
... (truncated)
Commits
67fad64 Merge pull request #6632 from skye:version
1246163 Merge pull request #6633 from zhangqiaorjc:jaxlib_rel
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.13.
Changelog
Sourced from jax's changelog.
... (truncated)
Commits
67fad64
Merge pull request #6632 from skye:version1246163
Merge pull request #6633 from zhangqiaorjc:jaxlib_rel24d126b
Update XLA.c7485b7
Bump jax version and changelog for jax 0.1.13 releasee29a9c4
[JAX] Bump xla_client._version and gate _ALLOW_ARG_PRUNING on the new version.850bd66
[JAX] Prune unused inputs in jit.e6bdcbb
Merge pull request #6625 from hawkinsp:xlaefe4a1f
Update XLA.ffecef5
Merge pull request #6611 from njunge94:pip-install-fork321ecbd
Merge pull request #6620 from hawkinsp:cudaDependabot 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)