rocm_jax/docs/jax.export.rst
George Necula cc73c50c41 [export] Improved the documentation.
In particular added the docstring for `Exported.call` method,
and fixed the formatting for `Exported.in_shardings_jax`.
2024-12-08 17:39:24 +01:00

58 lines
1.1 KiB
ReStructuredText

``jax.export`` module
=====================
.. automodule:: jax.export
:mod:`jax.export` is a library for exporting and serializing JAX functions
for persistent archival.
See the :ref:`export` documentation.
Classes
-------
.. autosummary::
:toctree: _autosummary
.. autoclass:: Exported
:members:
.. autoclass:: DisabledSafetyCheck
:members:
Functions
---------
.. autosummary::
:toctree: _autosummary
export
deserialize
minimum_supported_calling_convention_version
maximum_supported_calling_convention_version
default_export_platform
register_pytree_node_serialization
register_namedtuple_serialization
Functions related to shape polymorphism
---------------------------------------
.. autosummary::
:toctree: _autosummary
symbolic_shape
symbolic_args_specs
is_symbolic_dim
SymbolicScope
Constants
---------
.. data:: jax.export.minimum_supported_serialization_version
The minimum supported serialization version; see :ref:`export-calling-convention-version`.
.. data:: jax.export.maximum_supported_serialization_version
The maximum supported serialization version; see :ref:`export-calling-convention-version`.