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][source]#

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]]

Call a foreign function interface (FFI) target.

See theForeign function interface (FFI) tutorial for more information.

Likepure_callback(), the behavior offfi_call undervmap() depends on the value ofvmap_method. See thepure_callback() documentation for more details about the allowedvalues and examples of their behavior.

The current default behavior is to usevmap_method="sequential" whennot specified, but this behavior is deprecated, and in the future, thedefault will be to raise aNotImplementedError unlessvmap_method isexplicitly specified.

Parameters:

• target_name – the name of the XLA FFI custom call target that was registeredusingregister_ffi_target().

• result_shape_dtypes – an object, or sequence of objects, withshape anddtype attributes which are expected to match the shape and dtype ofthe custom call output or outputs.ShapeDtypeStruct is oftenused to define the elements ofresult_shape_dtypes.jax.core.abstract_token may be used to represent a token-typed output.

• has_side_effect – boolean specifying whether the custom call has sideeffects. WhenTrue, the FFI call will be executed even when theoutputs are not used.

• vmap_method – string specifying how the FFI call transforms undervmap() as described above.

• input_layouts – a sequence of layouts for each input argument. In each case,the layout can be (a)None indicating that this input is in defaultrow-major order, (b) aLayout specifying the axis order,or (c) a sequence of integers specifying the major-to-minor axisordering. Users who are familiar with XLA layouts should note that thisfunction expects layouts in major-to-minor order instead of theminor-to-major order that XLA uses. For example, a batch of row-majormatrices could be specified using the layout[0,1,2], whereas abatch of column-major matrices would have layout[0,2,1]. In bothof these examples, the leading/batch dimension is the “slowest” axis. Theinput_layouts parameter should be used to request the memory layoutexpected by the FFI call target, and XLA will ensure that the buffershave the correct layouts before the handler is executed.

• output_layouts – likeinput_layouts, but specifying the required layoutsfor the output arrays.

• input_output_aliases – a dictionary where the keys are input indices and thevalues are output indices. This mapping indicates which output arraysalias specific input arrays.

• custom_call_api_version – the version number of the custom call APIimplemented by the FFI targettarget_name. The only formallysupported version is the typed FFI API withcustom_call_api_version=4,but earlier unsupported custom calls can be executed using this argument.

• legacy_backend_config – for legacy targets implemented usingcustom_call_api_version<4, attributes are passed using the opaquestring representation provided by this argument. This parameter cannot beused withcustom_call_api_version>=4.

Returns:

A function that can be called with the input arrays as positional argumentsto execute the FFI handler. Any keyword arguments are passed as namedattributes to the FFI handler using XLA’s FFI interface.