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#
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 |