jaxlib.xla_client.register_custom_call_target

jaxlib.xla_client.register_custom_call_target(name, fn, platform='cpu', api_version=0, traits=<CustomCallTargetTraits.DEFAULT: 0>)

Registers a custom call target.

Parameters:

• name (str) – bytes containing the name of the function.

• fn (Any) – a PyCapsule object containing the function pointer.

• platform (str) – the target platform.

• api_version (int) – the XLA FFI version to use. Supported versions are: 0 for the untyped FFI and 1 for the typed FFI.

• traits (CustomCallTargetTraits) – custom call traits corresponding to XLA FFI handler traits.

Return type:

None