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.randommodulejax.experimental.serialize_executablemodulejax.experimental.sparsemodule
Contents