jax.extend.ffi.ffi_lowering

jax.extend.ffi.ffi_lowering#

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

为外部函数接口（FFI）目标构建一个降低规则。

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

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

参数:

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

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

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

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

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

返回类型:

mlir.LoweringRule

目录