Bump jax from 0.4.35 to 0.4.36

[dependabot]

Bumps jax from 0.4.35 to 0.4.36.

Release notes

Sourced from jax's releases.

JAX v0.4.36

• Breaking Changes

  • This release lands "stackless", an internal change to JAX's tracing machinery. We made trace dispatch purely a function of context rather than a function of both context and data. This let us delete a lot of machinery for managing data-dependent tracing: levels, sublevels, post_process_call, new_base_main, custom_bind, and so on. The change should only affect users that use JAX internals.

    If you do use JAX internals then you may need to update your code (see jax-ml/jax@c36e1f7 for clues about how to do this). There might also be version skew issues with JAX libraries that do this. If you find this change breaks your non-JAX-internals-using code then try the config.jax_data_dependent_tracing_fallback flag as a workaround, and if you need help updating your code then please file a bug.

  • jax.experimental.jax2tf.convert with native_serialization=False or with enable_xla=False have been deprecated since July 2024, with JAX version 0.4.31. Now we removed support for these use cases. jax2tf with native serialization will still be supported.

  • In jax.interpreters.xla, the xb, xc, and xe symbols have been removed after being deprecated in JAX v0.4.31. Instead use xb = jax.lib.xla_bridge, xc = jax.lib.xla_client, and xe = jax.lib.xla_extension.

  • The deprecated module jax.experimental.export has been removed. It was replaced by jax.export in JAX v0.4.30. See the migration guide for information on migrating to the new API.

  • The initial argument to jax.nn.softmax and jax.nn.log_softmax has been removed, after being deprecated in v0.4.27.

  • Calling np.asarray on typed PRNG keys (i.e. keys produced by jax.random.key) now raises an error. Previously, this returned a scalar object array.

  • The following deprecated methods and functions in jax.export have been removed:

    • jax.export.DisabledSafetyCheck.shape_assertions: it had no effect already.
    • jax.export.Exported.lowering_platforms: use platforms.
    • jax.export.Exported.mlir_module_serialization_version: use calling_convention_version.
    • jax.export.Exported.uses_shape_polymorphism: use uses_global_constants.
    • the lowering_platforms kwarg for jax.export.export: use platforms instead.

  • The kwargs symbolic_scope and symbolic_constraints from jax.export.symbolic_args_specs have been removed. They were deprecated in June 2024. Use scope and constraints instead.

  • Hashing of tracers, which has been deprecated since version 0.4.30, now results in a TypeError.

  • Refactor: JAX build CLI (build/build.py) now uses a subcommand structure and replaces previous build.py usage. Run python build/build.py --help for more details. Brief overview of the new subcommand options:

... (truncated)

Changelog

Sourced from jax's changelog.

jax 0.4.36

• Breaking Changes

  • This release lands "stackless", an internal change to JAX's tracing machinery. We made trace dispatch purely a function of context rather than a function of both context and data. This let us delete a lot of machinery for managing data-dependent tracing: levels, sublevels, post_process_call, new_base_main, custom_bind, and so on. The change should only affect users that use JAX internals.

    If you do use JAX internals then you may need to update your code (see jax-ml/jax@c36e1f7 for clues about how to do this). There might also be version skew issues with JAX libraries that do this. If you find this change breaks your non-JAX-internals-using code then try the config.jax_data_dependent_tracing_fallback flag as a workaround, and if you need help updating your code then please file a bug.

  • {func}jax.experimental.jax2tf.convert with native_serialization=False or with enable_xla=False have been deprecated since July 2024, with JAX version 0.4.31. Now we removed support for these use cases. jax2tf with native serialization will still be supported.

  • In jax.interpreters.xla, the xb, xc, and xe symbols have been removed after being deprecated in JAX v0.4.31. Instead use xb = jax.lib.xla_bridge, xc = jax.lib.xla_client, and xe = jax.lib.xla_extension.

  • The deprecated module jax.experimental.export has been removed. It was replaced by {mod}jax.export in JAX v0.4.30. See the migration guide for information on migrating to the new API.

  • The initial argument to {func}jax.nn.softmax and {func}jax.nn.log_softmax has been removed, after being deprecated in v0.4.27.

  • Calling np.asarray on typed PRNG keys (i.e. keys produced by :func:jax.random.key) now raises an error. Previously, this returned a scalar object array.

  • The following deprecated methods and functions in {mod}jax.export have been removed:

    • jax.export.DisabledSafetyCheck.shape_assertions: it had no effect already.
    • jax.export.Exported.lowering_platforms: use platforms.
    • jax.export.Exported.mlir_module_serialization_version: use calling_convention_version.
    • jax.export.Exported.uses_shape_polymorphism: use uses_global_constants.
    • the lowering_platforms kwarg for {func}jax.export.export: use platforms instead.

  • The kwargs symbolic_scope and symbolic_constraints from {func}jax.export.symbolic_args_specs have been removed. They were deprecated in June 2024. Use scope and constraints instead.

  • Hashing of tracers, which has been deprecated since version 0.4.30, now results in a TypeError.

  • Refactor: JAX build CLI (build/build.py) now uses a subcommand structure and replaces previous build.py usage. Run python build/build.py --help for

... (truncated)

Commits

• 259194a [Pallas] Fix shard_axis in dma_start interpret mode rule.
• 7e6620a JAX release 0.4.36.
• 23d5c10 [Mosaic:TPU] Fix fully replicated relayout
• 2a4a0e8 [jax:custom_partitioning] Implement SdyShardingRule to support
• f73fa7a Merge pull request #25290 from jakevdp:reduction-where
• a71f9a6 Merge pull request #25271 from jakevdp:fix-vector-norm
• e20a483 [JAX] Add end-to-end execution support in colocated Python API
• aaaee63 jnp.linalg.vector_norm: properly support multiple axes
• 29a8cce jax.numpy: require boolean dtype for where argument
• 3f5f3e1 [export] Removed __gpu$xla.gpu.triton (Pallas GPU) from the list of custom ca...
• 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 show <dependency name> 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)