JAX v0.3.17

v0.3.17

JAX release v0.3.17

Added 1
  • Add jax.pure_callback that enables calling back to pure Python functions from compiled functions
Fixed 1
  • Fix corner case issue in gradient of lax.pow with an exponent of zero
Removed 2
  • Remove jax.checkpoint concrete option following previous deprecation
  • Remove deprecated DeviceArray.tile() method
Deprecated 1
  • Deprecate DeviceArray.to_py() in favor of np.asarray(x)
  • GitHub commits.
  • Bugs
    • Fix corner case issue in gradient of lax.pow with an exponent of zero (#12041)
  • Breaking changes
    • jax.checkpoint, also known as jax.remat, no longer supports the concrete option, following the previous version's deprecation; see JEP 11830.
  • Changes
    • Added jax.pure_callback that enables calling back to pure Python functions from compiled functions (e.g. functions decorated with jax.jit or jax.pmap).
  • Deprecations:
    • The deprecated DeviceArray.tile() method has been removed. Use jax.numpy.tile (#11944).
    • DeviceArray.to_py() has been deprecated. Use np.asarray(x) instead.
View original

Upgraded? How did it go?

Discussion