jax.ffi.pycapsule

jax.ffi.pycapsule(funcptr)

將 ctypes 函數指標包裝在 PyCapsule 中。

此函數的主要用途，以及它存在於 jax.ffi 子模組中的原因，是為了包裝來自外部編譯函式庫的函數呼叫，以註冊為 XLA 自訂呼叫。

使用範例

import ctypes
import jax
from jax.lib import xla_client

libfoo = ctypes.cdll.LoadLibrary('./foo.so')
xla_client.register_custom_call_target(
    name="bar",
    fn=jax.ffi.pycapsule(libfoo.bar),
    platform=PLATFORM,
    api_version=API_VERSION
)

參數:

funcptr – 使用 ctypes 從動態函式庫載入的函數指標。

傳回:

包裝 funcptr 的不透明 PyCapsule 物件。