jax.ffi.ffi_lowering#
- jax.ffi.ffi_lowering(call_target_name, *, operand_layouts=None, result_layouts=None, backend_config=None, **lowering_args)[原始碼]#
為外部函數介面 (FFI) 目標建立降低規則。
預設情況下,此降低規則可以使用輸入和輸出抽象值來計算自訂呼叫的輸入和輸出類型及形狀,假設為列優先佈局。
請注意,傳遞給此函數作為元組的佈局應為次要到主要順序(如 XLA 所預期),而不是主要到次要順序,如
ffi_call()
和DeviceLocalLayout
所使用。如果關鍵字引數傳遞給降低規則,這些會被視為屬性,並新增到 backend_config。
- 參數:
call_target_name (
str
) – 自訂呼叫目標的名稱。operand_layouts (Sequence[FfiLayoutOptions] | None | None) – 每個運算元的佈局(維度順序)序列。預設情況下,運算元假定為列優先。
result_layouts (Sequence[FfiLayoutOptions] | None | None) – 每個結果的佈局(維度順序)序列。預設情況下,結果假定為列優先。
backend_config (Mapping[str, ir.Attribute] | str | None | None) – 自訂呼叫的配置資料。傳遞給降低規則的任何關鍵字引數都會新增到此字典。
lowering_args (Any) – 如果作為額外引數提供給此函數,則傳遞給
mlir.custom_call()
的任何其他引數也將會傳遞。
- 回傳類型:
mlir.LoweringRule