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[source]#
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 时使用 vmap_method="sequential"，但此行为已弃用，未来默认行为将是抛出 NotImplementedError，除非显式指定 vmap_method。

参数:

  • 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 处理程序。