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