jax.extend.ffi.register_ffi_target#

jax.extend.ffi.register_ffi_target(name, fn, platform='cpu', api_version=1, **kwargs)[source]#

Registers a foreign function target.

Parameters:

  • name (str) – the name of the target.

  • fn (Any) – a PyCapsule object containing the function pointer, or a dict where the keys are FFI stage names (e.g. “execute”) and the values are PyCapsule objects continaing a pointer to the handler for that stage.

  • platform (str) – the target platform.

  • api_version (int) – the XLA custom call API version to use. Supported versions are: 1 (default) for the typed FFI or 0 for the earlier “custom call” API.

  • kwargs (Any) – any extra keyword arguments are passed directly to register_custom_call_target() for more advanced use cases.

Return type:

None