jax.ffi.ffi_call#

jax.ffi.ffi_call(target_name: str, result_shape_dtypes: ResultMetadata, *, 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 | None | DeprecatedArg = DeprecatedArg()) → Callable[..., Array][source]#
jax.ffi.ffi_call(target_name: str, result_shape_dtypes: Sequence[ResultMetadata], *, 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 | None | DeprecatedArg = DeprecatedArg()) → Callable[..., 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 属性应与自定义调用输出的形状和数据类型匹配。ShapeDtypeStruct 通常用于定义 result_shape_dtypes 的元素。jax.core.abstract_token 可用于表示令牌类型的输出。

  • has_side_effect – 布尔值，指定自定义调用是否具有副作用。当为 True 时，即使不使用输出，也会执行 FFI 调用。

  • vmap_method – 字符串，指定如上所述 FFI 调用在 vmap() 下如何转换。

  • input_layouts – 每个输入参数的布局序列。在每种情况下，布局可以是 (a) None，表示此输入采用默认行主序；(b) 一个 Layout，指定轴顺序；或 (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 处理程序。