JAX release v0.3.16

• GitHub commits.
• Breaking changes
  • Support for NumPy 1.19 has been dropped, per the
    deprecation policy.
    Please upgrade to NumPy 1.20 or newer.
• Changes
  • Added jax.debug that includes utilities for runtime value debugging such at jax.debug.print and jax.debug.breakpoint.
  • Added new documentation for runtime value debugging
• Deprecations
  • jax.mask jax.shapecheck APIs have been removed. See #11557.
  • jax.experimental.loops has been removed. See #10278 for an alternative API.
  • jax.tree_util.tree_multimap has been removed. It has been deprecated since JAX release 0.3.5, and jax.tree_util.tree_map is a direct replacement.
  • Removed jax.experimental.stax; it has long been a deprecated alias of jax.example_libraries.stax.
  • Removed jax.experimental.optimizers; it has long been a deprecated alias of jax.example_libraries.optimizers.
  • jax.checkpoint, also known as jax.remat, has a new implementation switched on by default, meaning the old implementation is deprecated; see JEP 11830.