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][源代码]#
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 可用于表示 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 的版本号。唯一正式支持的版本是具有 custom_call_api_version=4 的类型化 FFI API，但可以使用此参数执行早期不受支持的自定义调用。

  • legacy_backend_config – 对于使用 custom_call_api_version<4 实现的旧版目标，属性通过此参数提供的非特定字符串表示形式传递。此参数不能与 custom_call_api_version>=4 一起使用。

返回:

一个函数，可以使用输入数组作为位置参数来调用，以执行 FFI 处理程序。任何关键字参数都将使用 XLA 的 FFI 接口作为命名属性传递给 FFI 处理程序。