Added {class}jax.numpy.ufunc, as well as {func}jax.numpy.frompyfunc, which can convert
any scalar-valued function into a {func}numpy.ufunc-like object, with methods such as
{meth}~jax.numpy.ufunc.outer, {meth}~jax.numpy.ufunc.reduce,
{meth}~jax.numpy.ufunc.accumulate, {meth}~jax.numpy.ufunc.at, and
{meth}~jax.numpy.ufunc.reduceat ({jax-issue}[#17054](https://github.com/google/jax/issues/17054)).
Added {func}jax.scipy.integrate.trapezoid.
When not running under IPython: when an exception is raised, JAX now filters out the
entirety of its internal frames from tracebacks. (Without the "unfiltered stack trace"
that previously appeared.) This should produce much friendlier-looking tracebacks. See
here for an example.
This behavior can be changed by setting JAX_TRACEBACK_FILTERING=remove_frames (for two
separate unfiltered/filtered tracebacks, which was the old behavior) or
JAX_TRACEBACK_FILTERING=off (for one unfiltered traceback).
jax2tf default serialization version is now 7, which introduces new shape
safety assertions.
Devices passed to jax.sharding.Mesh should be hashable. This specifically
applies to mock devices or user created devices. jax.devices() are
already hashable.
Breaking changes:
jax2tf now uses native serialization by default. See
the jax2tf documentation
for details and for mechanisms to override the default.
The option --jax_coordination_service has been removed. It is now always
True.
jax.jaxpr_util has been removed from the public JAX namespace.
JAX_USE_PJRT_C_API_ON_TPU no longer has an effect (i.e. it always defaults to true).
The backwards compatibility flag --jax_host_callback_ad_transforms
introduced in December 2021, has been removed.
Deprecations:
Several jax.numpy APIs have been deprecated following
NumPy NEP-52:
jax.numpy.NINF has been deprecated. Use -jax.numpy.inf instead.
jax.numpy.PZERO has been deprecated. Use 0.0 instead.
jax.numpy.NZERO has been deprecated. Use -0.0 instead.
jax.numpy.issubsctype(x, t) has been deprecated. Use jax.numpy.issubdtype(x.dtype, t).
jax.numpy.row_stack has been deprecated. Use jax.numpy.vstack instead.
jax.numpy.in1d has been deprecated. Use jax.numpy.isin instead.
jax.numpy.trapz has been deprecated. Use jax.scipy.integrate.trapezoid instead.
jax.scipy.linalg.tril and jax.scipy.linalg.triu have been deprecated,
following SciPy. Use jax.numpy.tril and jax.numpy.triu instead.
jax.lax.prod has been removed after being deprecated in JAX v0.4.11.
Use the built-in math.prod instead.
A number of exports from jax.interpreters.xla related to defining
HLO lowering rules for custom JAX primitives have been deprecated. Custom
primitives should be defined using the StableHLO lowering utilities in
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 show ignore conditions` will show all of the ignore conditions of the specified dependency
- `@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)
Updates the requirements on jax[cpu] to permit the latest version.
Changelog
Sourced from jax[cpu]'s changelog.
... (truncated)
Commits
d477b92
Start jax and jaxlib 0.4.16 releasefe0f12e
[Pallas] Wire up cost estimates in Mosaic paramsee846fa
Merge pull request #17642 from apaszke:fix-no-teardown9fa96f8
Raise SkipTest before running setUp32ee27b
[callbacks] Add support for shardable ordered effects.152af70
Update XLA dependency to use revisionbcc545a
Update XLA dependency to use revisionf863cfb
Relax some test tolerances to fix failures on Linux aarch64.bf40f75
Update XLA dependency to use revisione78d8a3
[Pallas] Add Mosaic lowering rule for fpowi.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 show