https://github.com/google/jax
Raw File
Tip revision: 03ed46dcd6dacbaf6d8d3bc2d86a4d15a604be66 authored by Yash Katariya on 29 December 2022, 21:00:06 UTC
Remove `__dict__` from C++ Array. Also disallow setting attributes on Array.
Tip revision: 03ed46d
jax.experimental.rst
.. currentmodule:: jax.experimental

jax.experimental package
========================

``jax.experimental.optix`` has been moved into its own Python package
(https://github.com/deepmind/optax).

``jax.experimental.ann`` has been moved into ``jax.lax``.

Experimental Modules
--------------------

.. toctree::
    :maxdepth: 1

    jax.experimental.checkify
    jax.experimental.global_device_array
    jax.experimental.host_callback
    jax.experimental.maps
    jax.experimental.pjit
    jax.experimental.sparse
    jax.experimental.jet

Experimental APIs
-----------------

.. autosummary::
   :toctree: _autosummary

   enable_x64
   disable_x64

   jax.experimental.checkify.checkify
   jax.experimental.checkify.check
   jax.experimental.checkify.check_error
back to top