Toggle navigation

.rst .pdf

Contents

Warning

This page was created from a pull request (#9655).

jax.experimental package

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

Experimental Modules

• jax.experimental.ann module
• jax.experimental.global_device_array module
• jax.experimental.host_callback module
• jax.experimental.loops module
• jax.experimental.maps module
• jax.experimental.pjit module
• jax.experimental.sparse 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(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.

previous

jax.example_libraries.stax module

next

jax.experimental.ann module