jax-ml/jax jax-v0.4.28

JAX v0.4.28

on GitHub

Bug fixes

• Reverted a change to make_jaxpr that was breaking Equinox (#21116).

Deprecations & removals

• The kind argument to jax.numpy.sort and jax.numpy.argsort
  is now removed. Use stable=True or stable=False instead.
• Removed get_compute_capability from the jax.experimental.pallas.gpu
  module. Use the compute_capability attribute of a GPU device, returned
  by jax.devices or jax.local_devices, instead.

Changes

• The minimum jaxlib version of this release is 0.4.27.