jax.experimental module
Contents
jax.experimental module#
jax.experimental.optix has been moved into its own Python package(deepmind/optax).
jax.experimental.ann has been moved intojax.lax.
Experimental Modules#
jax.experimental.checkifymodulejax.experimental.compilation_cachemodulejax.experimental.custom_dcemodulejax.experimental.custom_partitioningmodulejax.experimental.jetmodulejax.experimental.key_reusemodulejax.experimental.mesh_utilsmodulejax.experimental.multihost_utilsmodulejax.experimental.pallasmodulejax.experimental.serialize_executablemodulejax.experimental.sparsemodule
Experimental APIs#
| Experimental context manager to temporarily enable X64 mode. |
Experimental context manager to temporarily disable X64 mode. |
Contents