jax.ffi.register_ffi_target#

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

註冊外部函數目標。

參數:

  • name (str) – 目標的名稱。

  • fn (Any) – PyCapsule 物件，包含函數指標；或 dict，其中鍵為 FFI 階段名稱 (例如 “execute”)，值為 PyCapsule 物件，包含指向該階段處理常式的指標。

  • platform (str) – 目標平台。

  • api_version (int) – 要使用的 XLA 自訂呼叫 API 版本。支援的版本為：1 (預設) 用於型別化 FFI，或 0 用於較早的「自訂呼叫」API。

  • kwargs (Any) – 任何額外的關鍵字引數都會直接傳遞給 register_custom_call_target() 以用於更進階的用例。

回傳型別:

None