JAX v0.4.38

v0.4.38

JAX v0.4.38

Added 2
  • Add jax.tree.flatten_with_path and jax.tree.map_with_path as shortcuts of the corresponding tree_util functions
  • Add jax.lax.split as a primitive version of jax.numpy.split that yields a more compact transpose during automatic differentiation
Changed 1
  • jax.export.export can be used for device-polymorphic export with shardings constructed with jax.sharding.AbstractMesh
Removed 4
  • Remove jax.core.check_eqn, jax.core.check_type, jax.core.check_valid_jaxtype, and jax.core.non_negative_dim
  • Remove jax.lib.xla_bridge.xla_client and jax.lib.xla_bridge.default_backend
  • Remove jax.lib.xla_client._xla and jax.lib.xla_client.bfloat16
  • Remove jax.numpy.round_
Deprecated 1
  • Deprecate a number of APIs in the internal jax.core namespace

  • Changes:

    • jax.tree.flatten_with_path and jax.tree.map_with_path are added as shortcuts of the corresponding tree_util functions.

  • Deprecations

    • a number of APIs in the internal jax.core namespace have been deprecated. Most were no-ops, were little-used, or can be replaced by APIs of the same name in jax.extend.core; see the documentation for {mod}jax.extend for information on the compatibility guarantees of these semi-public extensions.
    • Several previously-deprecated APIs have been removed, including:
      • from jax.core: check_eqn, check_type, check_valid_jaxtype, and non_negative_dim.
      • from jax.lib.xla_bridge: xla_client and default_backend.
      • from jax.lib.xla_client: _xla and bfloat16.
      • from jax.numpy: round_.

  • New Features

    • jax.export.export can be used for device-polymorphic export with shardings constructed with {func}jax.sharding.AbstractMesh. See the jax.export documentation.
    • Added jax.lax.split. This is a primitive version of jax.numpy.split, added because it yields a more compact transpose during automatic differentiation.
View original

Upgraded? How did it go?

Discussion