jax.ffi.ffi_lowering#

jax.ffi.ffi_lowering(call_target_name, *, operand_layouts=None, result_layouts=None, backend_config=None, skip_ffi_layout_processing=False, **lowering_args)[源代码]#

为外部函数接口 (FFI) 目标构建降级规则。

默认情况下，此降级规则可以使用输入和输出抽象值来计算自定义调用的输入和输出类型与形状，假设为行主序布局。

请注意，传递给此函数的布局元组应采用次到主顺序（XLA 所需），而不是 ffi_call() 和 Layout 所使用的主到次顺序。

如果关键字参数传递给降级规则，这些将被视为属性，并添加到 backend_config。

参数:

  • call_target_name (str) – 自定义调用目标的名称。

  • operand_layouts (Sequence[FfiLayoutOptions] | None) – 每个操作数的布局（维度顺序）序列。默认情况下，假定操作数为行主序。

  • result_layouts (Sequence[FfiLayoutOptions] | None) – 每个结果的布局（维度顺序）序列。默认情况下，假定结果为行主序。

  • backend_config (Mapping[str, ir.Attribute] | str | None) – 自定义调用的配置数据。传递给降级规则的任何关键字参数都将添加到此字典中。

  • lowering_args (Any) – 如果作为此函数的额外参数提供，mlir.custom_call() 的任何其他参数也将被传递。

  • skip_ffi_layout_processing (bool) – 如果为 true，则跳过传递给降级规则的操作数和结果布局参数的处理。

返回类型:

mlir.LoweringRule