jax.extend.ffi.ffi_call#

jax.extend.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.extend.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]

调用外部函数接口 (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 可用于表示令牌类型的输出。

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