Skip to content

Plugins and runtime

Loading a plugin and reading the state of the runtime.

initialize 

initialize(plugin: Plugin | Path | None = None, *, platform: str | None = None, options: Mapping[str, OptionValue] | None = None, jax_targets: bool = True) -> RuntimeInfo

Loads a PJRT plugin and returns details of its runtime.

Calling it is optional, as the first function that needs a plugin loads the default one. It is for choosing the plugin or its options, and has to come before that first use.

Without a plugin, find_plugin picks one for platform. A path is loaded with the default options of platform, which is cpu if not given. options add to or replace the plugin's client options. A process loads one plugin, so later calls return the details of the first.

With jax_targets the FFI handlers that jaxlib ships are registered with the plugin, which jnp.linalg programs need, and so are the targets registered earlier with register_ffi_target. Pass False to register none.

available_plugins 

available_plugins() -> tuple[Plugin, ...]

Returns every installed plugin, in the order of the module docstring.

find_plugin 

find_plugin(platform: str | None = None) -> Plugin

Returns the plugin to load for a platform.

The path in CYXLA_PJRT_PLUGIN wins over installed plugins. Without a platform, the one in CYXLA_PLATFORM is used, and without that, the first installed platform in PLATFORM_PREFERENCE.

Plugin

Bases: NamedTuple

A PJRT plugin library and the options to create its client with.

platform instance-attribute 

platform: str

Platform that the plugin runs programs on, which is also the JAX lowering platform.

path instance-attribute 

path: Path

Path of the plugin library.

source instance-attribute 

source: str

Where the plugin was found, such as the distribution that installed it.

create_options class-attribute instance-attribute 

create_options: Mapping[str, OptionValue] = MappingProxyType({})

Options of PJRT_Client_Create, which initialize merges with the caller's.

preload class-attribute instance-attribute 

preload: tuple[Path, ...] = ()

Libraries to load before the plugin, so the dynamic loader finds them by soname.

PluginModule

Bases: Protocol

A module that a cyxla.plugins entry point names.

PLATFORM instance-attribute 

PLATFORM: str

Platform that the plugin runs programs on.

CREATE_OPTIONS instance-attribute 

CREATE_OPTIONS: Mapping[str, OptionValue]

Options of PJRT_Client_Create for this plugin.

library_path 

library_path() -> Path

Returns the path of the plugin library.

RuntimeInfo

Bases: NamedTuple

Details of the PJRT runtime.

platform instance-attribute 

platform: str

Platform name that the plugin reports, such as cpu or cuda.

aliased_inputs instance-attribute 

aliased_inputs: bool

Whether input buffers share host memory, so filling an input needs no copy.

inplace_outputs instance-attribute 

inplace_outputs: bool

Whether outputs are read from the buffers' memory instead of copied to the host.

runtime_info 

runtime_info() -> RuntimeInfo

Returns the PJRT API version, the platform and whether executions avoid copies.

is_initialized 

is_initialized() -> bool

Returns whether a PJRT plugin is loaded.

ErrorInfo

Bases: NamedTuple

PJRT errors recorded since the runtime was initialized.

error_info 

error_info() -> ErrorInfo

Returns the number of PJRT errors so far and the last message.

OptionValue module-attribute 

OptionValue = bool | int | float | str | tuple[int, ...]

A value of a PJRT_Client_Create option, one of the types of PJRT_NamedValue.

PLUGIN_ENVIRONMENT_VARIABLE module-attribute 

PLUGIN_ENVIRONMENT_VARIABLE = 'CYXLA_PJRT_PLUGIN'

Environment variable with the path of a plugin that overrides the installed ones.

PLATFORM_ENVIRONMENT_VARIABLE module-attribute 

PLATFORM_ENVIRONMENT_VARIABLE = 'CYXLA_PLATFORM'

Environment variable with the platform to load when none is passed.