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#


Experimental context manager to temporarily enable X64 mode.


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.


Raise an Exception if error represents a failure.