jax.extend.ffi.register_ffi_target#

jax.extend.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