jax.experimental package#

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

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

Experimental Modules#

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(fun[, 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.