jax-ml/jax jax-v0.4.38

JAX v0.4.38

on GitHub

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.