rocm_jax/docs/jax.lib.rst

``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
  get_backend
  get_compile_options

jax.lib.xla_client
------------------

.. currentmodule:: jaxlib.xla_client

.. autosummary::
  :toctree: _autosummary

  register_custom_call_target
DOC: make API doc titles more uniform 2023-01-18 10:59:42 -08:00			``jax.lib`` module
			`==================`
DOC: fix docstring and add docs. 2021-09-27 09:48:27 -07:00			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`
			`get_backend`
			`get_compile_options`

			`jax.lib.xla_client`
			`------------------`

			`.. currentmodule:: jaxlib.xla_client`

			`.. autosummary::`
Add `register_custom_call_target` to `xla_client` API docs. This function is (for better or worse) user facing for custom call users. I think it's worth having this in the API docs. 2024-06-27 13:50:11 -04:00			`:toctree: _autosummary`

			`register_custom_call_target`