Custom calls¶

Registers an FFI handler under a name, for programs that call that target.

The signature follows jax.ffi.register_ffi_target, so a library can pass the same arguments to both. Registering twice with the same handler does nothing, and with another handler raises an error, as the plugin keeps a registration for the whole process. A handler for a platform other than the plugin's waits until a plugin for that platform loads. The argument is typed Any for the reason given at handler_address.

Registers the handlers that JAX knows, and returns how many were new.

This covers the LAPACK kernels of jaxlib, which make jnp.linalg programs compile on the CPU, and everything else in the registry of JAX, such as the handlers of libraries that registered with jax.ffi.register_ffi_target before cyxla was imported. It also loads the LAPACK functions of SciPy that the kernels call. Handlers of the older untyped interface are left out, as the plugin cannot run them, and so are handlers that the plugin cannot take, which raise a warning. A plugin without FFI support gets none.

Registers every typed handler that libraries register with JAX from now on.

A library that supports JAX's FFI calls jax.ffi.register_ffi_target when it is imported. After this call, such a handler is also registered with the PJRT plugin, so the library works in programs that cyxla compiles without changes. Handlers of the older untyped interface are not covered. Registrations from before the call are not seen here, but loading the plugin registers them from the registry of JAX. Importing cyxla calls this unless CYXLA_MIRROR_JAX is 0, and calling it again does nothing.