Bump jax from 0.2.8 to 0.3.6

[dependabot[bot]]

Bumps jax from 0.2.8 to 0.3.6.

Release notes

Sourced from jax's releases.

JAX release v0.3.6

• Changes:
  • Upgraded libtpu wheel to the fixed version. Fixes #10218.

JAX release v0.3.5

Changes

• added jax.random.loggamma & improved behavior of jax.random.beta and jax.random.dirichlet for small parameter values (#9906).
• the private lax_numpy submodule is no longer exposed in the jax.numpy namespace (#10029).
• added array creation routines jax.numpy.frombuffer, jax.numpy.fromfunction, and jax.numpy.fromstring (#10049).
• DeviceArray.copy() now returns a DeviceArray rather than a np.ndarray (#10069)
• added jax.scipy.linalg.rsf2csf
• Deprecations:
  • jax.nn.normalize is being deprecated. Use jax.nn.standardize instead (#9899).
  • jax.tree_util.tree_multimap is deprecated. Use jax.tree_util.tree_map instead (#5746).
  • jax.experimental.sharded_jit is deprecated. Use pjit instead.

JAX release v0.3.4

Fix a bug introduced in #9923.

JAX release v0.3.3

• GitHub commits.

Jax release v0.3.1

• Changes:
  • 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):
      ...
    

  • Added jax.scipy.linalg.schur, jax.scipy.linalg.sqrtm, jax.scipy.signal.csd, jax.scipy.signal.stft, jax.scipy.signal.welch.

Jax release v0.3.0

• Changes
  • jax version has been bumped to 0.3.0. Please see the design doc for the explanation.

Jaxlib release v0.3.0

• Changes

... (truncated)

Changelog

Sourced from jax's changelog.

jaxlib 0.3.6 (Unreleased)

jax 0.3.7 (Unreleased)

• GitHub commits.

jax 0.3.6 (April 12, 2022)

• GitHub commits.
• Changes:
  • Upgraded libtpu wheel to the fixed version. Fixes #10218.

jax 0.3.5 (April 7, 2022)

• GitHub commits.
• Changes:
  • added {func}jax.random.loggamma & improved behavior of {func}jax.random.beta and {func}jax.random.dirichlet for small parameter values ({jax-issue}[#9906](https://github.com/google/jax/issues/9906)).
  • the private lax_numpy submodule is no longer exposed in the jax.numpy namespace ({jax-issue}[#10029](https://github.com/google/jax/issues/10029)).
  • added array creation routines {func}jax.numpy.frombuffer, {func}jax.numpy.fromfunction, and {func}jax.numpy.fromstring ({jax-issue}[#10049](https://github.com/google/jax/issues/10049)).
  • DeviceArray.copy() now returns a DeviceArray rather than a np.ndarray ({jax-issue}[#10069](https://github.com/google/jax/issues/10069))
  • added {func}jax.scipy.linalg.rsf2csf
• Deprecations:
  • {func}jax.nn.normalize is being deprecated. Use {func}jax.nn.standardize instead ({jax-issue}[#9899](https://github.com/google/jax/issues/9899)).
  • {func}jax.tree_util.tree_multimap is deprecated. Use {func}jax.tree_util.tree_map instead ({jax-issue}[#5746](https://github.com/google/jax/issues/5746)).
  • jax.experimental.sharded_jit is deprecated. Use pjit instead.

jaxlib 0.3.5 (April 7, 2022)

• Bug fixes
  • Fixed a bug where double-precision complex-to-real IRFFTs would mutate their input buffers on GPU ({jax-issue}[#9946](https://github.com/google/jax/issues/9946)).
  • Fixed incorrect constant-folding of complex scatters ({jax-issue}[#10159](https://github.com/google/jax/issues/10159))

jax 0.3.4 (March 18, 2022)

• GitHub commits.

jax 0.3.3 (March 17, 2022)

• GitHub commits.

jax 0.3.2 (March 16, 2022)

• GitHub commits.
• Changes:
  • The functions jax.ops.index_update, jax.ops.index_add, which were

... (truncated)

Commits

• 6ba9fb6 Upgrade the bazel version to 5.1.1
• c06eff8 Merge pull request #10245 from google:yashk2810-patch-7
• 5fd78ea Bump the libtpu version to prepare for JAX release
• 9455254 [MHLO] Add a direct MHLO lowering for pjit_p, which lacked one.
• 3136004 Fix the pytype error. PyType is looking for a init method. This does not ...
• a2c2d9a [JAX] Adds the approx_top_k_p bridge.
• 7be37ab Merge pull request #10120 from mattjj:djax-latest
• 4354f35 prototyping dynamic shapes
• fb6a143 Merge pull request #9723 from sharadmv:jaxpr-effects
• b051be4 [MHLO] Switch sharded_jit dispatch path to use MHLO lowering.
• Additional commits viewable in compare view

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)

[dependabot[bot]]

Superseded by #290.