jax.test_util.JaxTestCase and jax.test_util.JaxTestLoader are now deprecated.
The suggested replacement is to use parametrized.TestCase directly. For tests that
rely on custom asserts such as JaxTestCase.assertAllClose(), the suggested replacement
is to use standard numpy testing utilities such as numpy.testing.assert_allclose(),
which work directly with JAX arrays (#9620 ).
jax.test_util.JaxTestCase now sets jax_numpy_rank_promotion='raise' by default
(#9562 ). To recover the previous behavior, use the new
jax.test_util.with_config decorator:
@jtu.with_config(jax_numpy_rank_promotion='allow')
class MyTestCase(jtu.JaxTestCase):
...
Support for NumPy 1.18 has been dropped, per the [deprecation policy](https://jax.readthedocs.io/en/latest/ deprecation.html). Please upgrade to a supported NumPy version.
The host_callback primitives have been simplified to drop the special autodiff handling for hcb.id_tap and id_print. From now on, only the primals are tapped. The old behavior can be obtained (for a limited time) by setting the JAX_HOST_CALLBACK_AD_TRANSFORMS environment variable, or the --flax_host_callback_ad_transforms flag. Additionally, added documentation for how to implement the old behavior using JAX custom AD APIs ({jax-issue}[#8678](https://github.com/google/jax/issues/8678)).
Sorting now matches the behavior of NumPy for 0.0 and NaN regardless of the bit representation. In particular, 0.0 and -0.0 are now treated as equivalent, where previously -0.0 was treated as less than 0.0. Additionally all NaN representations are now treated as equivalent and sorted to the end of the array. Previously negative NaN values were sorted to the front of the array, and NaN values with different internal bit representations were not treated as equivalent, and were sorted according to those bit patterns ({jax- issue}[#9178](https://github.com/google/jax/issues/9178)).
{func}jax.numpy.unique now treats NaN values in the same way as np.unique in NumPy versions 1.21 and newer: at most one NaN value will appear in the uniquified output ({jax-issue}9184).
The functions jax.ops.index_update, jax.ops.index_add, which were
deprecated in 0.2.22, have been removed. Please use
the .at property on JAX arrays
instead, e.g., x.at[idx].set(y).
Moved jax.experimental.ann.approx_*_k into jax.lax. These functions are
optimized alternatives to jax.lax.top_k.
{func}jax.numpy.broadcast_arrays and {func}jax.numpy.broadcast_to now require scalar
or array-like inputs, and will fail if they are passed lists (part of {jax-issue}[#7737](https://github.com/google/jax/issues/7737)).
The standard jax[tpu] install can now be used with Cloud TPU v4 VMs.
pjit now works on CPU (in addition to previous TPU and GPU support).
jaxlib 0.3.2 (March 16, 2022)
Changes
XlaComputation.as_hlo_text() now supports printing large constants by
passing boolean flag print_large_constants=True.
Deprecations:
The .block_host_until_ready() method on JAX arrays has been deprecated.
Use .block_until_ready() instead.
jax.test_util.JaxTestCase and jax.test_util.JaxTestLoader are now deprecated.
The suggested replacement is to use parametrized.TestCase directly. For tests that
rely on custom asserts such as JaxTestCase.assertAllClose(), the suggested replacement
is to use standard numpy testing utilities such as {func}numpy.testing.assert_allclose(),
which work directly with JAX arrays ({jax-issue}[#9620](https://github.com/google/jax/issues/9620)).
jax.test_util.JaxTestCase now sets jax_numpy_rank_promotion='raise' by default
({jax-issue}[#9562](https://github.com/google/jax/issues/9562)). To recover the previous behavior, use the new
jax.test_util.with_config decorator:
@jtu.with_config(jax_numpy_rank_promotion='allow')
class MyTestCase(jtu.JaxTestCase):
...
... (truncated)
Commits
5d7f639 Add small and big matmul to api_benchmarks.
53f52cb Merge pull request #9942 from jakevdp:glu-fix
ae631e9 Merge pull request #9944 from jakevdp:poly-refactor
603bb3c lax_numpy: move poly functions into numpy.polynomial
2d79a64 Merge pull request #9943 from jakevdp:ufunc-imports
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.3.3.
Release notes
Sourced from jax's releases.
... (truncated)
Changelog
Sourced from jax's changelog.
... (truncated)
Commits
5d7f639
Add small and big matmul to api_benchmarks.53f52cb
Merge pull request #9942 from jakevdp:glu-fixae631e9
Merge pull request #9944 from jakevdp:poly-refactor603bb3c
lax_numpy: move poly functions into numpy.polynomial2d79a64
Merge pull request #9943 from jakevdp:ufunc-imports131c005
lax_numpy refactoring: directly import ufuncsd7bd93f
Merge pull request #9925 from jakevdp:refactor-index-opsc762e07
jax.nn.glu: fix static argname issuec3a4a6e
Revert previous change0a72adb
lax_numpy: factor out indexing tricksDependabot 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)