jax.ffi.ffi_call#

jax.ffi.ffi_call(target_name: str, result_shape_dtypes: ResultMetadata, *deprecated_args: ArrayLike, has_side_effect: bool = False, vmap_method: str | None = None, input_layouts: Sequence[FfiLayoutOptions] | None = None, output_layouts: FfiLayoutOptions | Sequence[FfiLayoutOptions] | None = None, input_output_aliases: dict[int, int] | None = None, custom_call_api_version: int = 4, legacy_backend_config: str | None = None, vectorized: bool | DeprecatedArg = Deprecated, **deprecated_kwargs: Any) → Callable[..., Array] | Array[原始碼]#
jax.ffi.ffi_call(target_name: str, result_shape_dtypes: Sequence[ResultMetadata], *deprecated_args: ArrayLike, has_side_effect: bool = False, vmap_method: str | None = None, input_layouts: Sequence[FfiLayoutOptions] | None = None, output_layouts: FfiLayoutOptions | Sequence[FfiLayoutOptions] | None = None, input_output_aliases: dict[int, int] | None = None, custom_call_api_version: int = 4, legacy_backend_config: str | None = None, vectorized: bool | DeprecatedArg = Deprecated, **deprecated_kwargs: Any) → Callable[..., Sequence[Array]] | Sequence[Array]

呼叫外部函數介面 (FFI) 目標。

請參閱外部函數介面 (FFI) 教學以取得更多資訊。

如同 pure_callback()，ffi_call 在 vmap() 下的行為取決於 vmap_method 的值。請參閱 pure_callback() 文件以取得關於允許值及其行為範例的更多詳細資訊。

目前的預設行為是在未指定時使用 vmap_method="sequential"，但此行為已被棄用，未來，除非明確指定 vmap_method，否則預設行為將會引發 NotImplementedError。

參數：

  • target_name – 使用 register_ffi_target() 註冊的 XLA FFI 自訂呼叫目標名稱。

  • result_shape_dtypes – 一個物件或物件序列，具有 shape 和 dtype 屬性，預期與自訂呼叫輸出或輸出的形狀和 dtype 相符。ShapeDtypeStruct 通常用於定義 result_shape_dtypes 的元素。jax.core.abstract_token 可用於表示 token 型別的輸出。

  • has_side_effect – 布林值，指定自訂呼叫是否具有副作用。當 True 時，即使未使用輸出，仍會執行 FFI 呼叫。

  • vmap_method – 字串，指定 FFI 呼叫如何在 vmap() 下轉換，如上所述。

  • input_layouts – 每個輸入引數的版面配置序列。在每種情況下，版面配置可以是 (a) None，表示此輸入為預設的列優先順序，(b) DeviceLocalLayout，指定軸順序，或 (c) 整數序列，指定主軸到次軸的軸順序。熟悉 XLA 版面配置的使用者應注意，此函數預期版面配置採用主軸到次軸的順序，而不是 XLA 使用的次軸到主軸的順序。例如，可以使用版面配置 [0, 1, 2] 指定一批列優先矩陣，而一批行優先矩陣將具有版面配置 [0, 2, 1]。在這些範例中，前導/批次維度都是「最慢」的軸。input_layouts 參數應用於請求 FFI 呼叫目標預期的記憶體版面配置，而 XLA 將確保緩衝區在處理常式執行之前具有正確的版面配置。

  • output_layouts – 類似於 input_layouts，但指定輸出陣列的必要版面配置。

  • input_output_aliases – 字典，其中鍵是輸入索引，值是輸出索引。此映射指示哪些輸出陣列是特定輸入陣列的別名。

  • custom_call_api_version – FFI 目標 target_name 實作的自訂呼叫 API 版本號碼。唯一正式支援的版本是型別化的 FFI API，其 custom_call_api_version=4，但可以使用此引數執行較早的不支援自訂呼叫。

  • legacy_backend_config – 對於使用 custom_call_api_version<4 實作的舊版目標，屬性會使用此引數提供的不透明字串表示法傳遞。此參數不能與 custom_call_api_version>=4 一起使用。

傳回：

一個函數，可以使用輸入陣列作為位置引數來呼叫，以執行 FFI 處理常式。任何關鍵字引數都會使用 XLA 的 FFI 介面作為具名屬性傳遞至 FFI 處理常式。