rocm_jax/docs/jax.lib.rst
2023-01-18 10:59:42 -08:00

38 lines
651 B
ReStructuredText

``jax.lib`` module
==================
The `jax.lib` package is a set of internal tools and types for bridging between
JAX's Python frontend and its XLA backend.
jax.lib.xla_bridge
------------------
.. currentmodule:: jax.lib.xla_bridge
.. autosummary::
:toctree: _autosummary
default_backend
device_count
get_backend
get_compile_options
local_device_count
process_index
jax.lib.xla_client
------------------
.. currentmodule:: jaxlib.xla_client
.. autosummary::
:toctree: _autosummary
jax.lib.xla_extension
---------------------
.. currentmodule:: jaxlib.xla_extension
.. autosummary::
:toctree: _autosummary
Device