jax.experimental module

jax.experimental.optix has been moved into its own Python package (deepmind/optax).

jax.experimental.ann has been moved into jax.lax.

Experimental Modules

• jax.experimental.checkify module
• jax.experimental.host_callback module
• jax.experimental.maps module
• jax.experimental.pjit module
• jax.experimental.sparse module
• jax.experimental.jet module
• jax.experimental.custom_partitioning module

Experimental APIs

enable_x64([new_val])	Experimental context manager to temporarily enable X64 mode.
disable_x64()	Experimental context manager to temporarily disable X64 mode.
jax.experimental.checkify.checkify(f[, errors])	Functionalize check calls in fun, and optionally add run-time error checks.
jax.experimental.checkify.check(pred, msg, ...)	Check a predicate, add an error with msg if predicate is False.
jax.experimental.checkify.check_error(error)	Raise an Exception if error represents a failure.