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.array_api
modulejax.experimental.checkify
modulejax.experimental.host_callback
modulejax.experimental.maps
modulejax.experimental.pjit
modulejax.experimental.sparse
modulejax.experimental.jet
modulejax.experimental.custom_partitioning
modulejax.experimental.multihost_utils
modulejax.experimental.compilation_cache
modulejax.experimental.key_reuse
modulejax.experimental.mesh_utils
module
Experimental APIs#
|
Experimental context manager to temporarily enable X64 mode. |
Experimental context manager to temporarily disable X64 mode. |
|
|
Functionalize check calls in fun, and optionally add run-time error checks. |
|
Check a predicate, add an error with msg if predicate is False. |
Raise an Exception if |