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>)

註冊自訂呼叫目標。

參數:

• name (str) – 包含函式名稱的位元組。

• fn (Any) – 包含函式指標的 PyCapsule 物件。

• platform (str) – 目標平台。

• api_version (int) – 要使用的 XLA FFI 版本。支援的版本有：0 代表未型別化 FFI，以及 1 代表型別化 FFI。

• traits (CustomCallTargetTraits) – 對應於 XLA FFI 處理常式特徵的自訂呼叫特徵。

返回類型:

無