2022-02-04 14:40:14 -08:00
|
|
|
.. currentmodule:: jax.experimental
|
|
|
|
|
|
2023-01-18 10:59:42 -08:00
|
|
|
``jax.experimental`` module
|
|
|
|
|
===========================
|
2019-01-15 20:14:19 -05:00
|
|
|
|
2020-11-19 08:02:46 -08:00
|
|
|
``jax.experimental.optix`` has been moved into its own Python package
|
2022-02-04 14:40:14 -08:00
|
|
|
(https://github.com/deepmind/optax).
|
|
|
|
|
|
2022-03-01 14:46:04 -08:00
|
|
|
``jax.experimental.ann`` has been moved into ``jax.lax``.
|
|
|
|
|
|
2022-02-04 14:40:14 -08:00
|
|
|
Experimental Modules
|
|
|
|
|
--------------------
|
2020-11-19 08:02:46 -08:00
|
|
|
|
2019-01-15 20:14:19 -05:00
|
|
|
.. toctree::
|
|
|
|
|
:maxdepth: 1
|
|
|
|
|
|
2022-07-26 14:47:36 -07:00
|
|
|
jax.experimental.checkify
|
2023-11-19 11:18:43 -08:00
|
|
|
jax.experimental.compilation_cache
|
2025-01-27 11:16:00 -05:00
|
|
|
jax.experimental.custom_dce
|
2024-10-08 11:13:04 +01:00
|
|
|
jax.experimental.custom_partitioning
|
|
|
|
|
jax.experimental.jet
|
2023-12-11 12:03:48 -08:00
|
|
|
jax.experimental.key_reuse
|
2024-02-13 00:56:21 +00:00
|
|
|
jax.experimental.mesh_utils
|
2024-10-08 11:13:04 +01:00
|
|
|
jax.experimental.multihost_utils
|
|
|
|
|
jax.experimental.pallas
|
|
|
|
|
jax.experimental.pjit
|
2024-05-14 13:36:35 -07:00
|
|
|
jax.experimental.serialize_executable
|
2024-05-13 13:04:15 -07:00
|
|
|
jax.experimental.shard_map
|
2024-10-08 11:13:04 +01:00
|
|
|
jax.experimental.sparse
|
2019-01-15 20:14:19 -05:00
|
|
|
|
2022-02-04 14:40:14 -08:00
|
|
|
Experimental APIs
|
|
|
|
|
-----------------
|
|
|
|
|
|
|
|
|
|
.. autosummary::
|
|
|
|
|
:toctree: _autosummary
|
2021-01-25 13:23:15 -08:00
|
|
|
|
2022-02-04 14:40:14 -08:00
|
|
|
enable_x64
|
|
|
|
|
disable_x64
|
