jax.extend.ffi.register_ffi_target

jax.extend.ffi.register_ffi_target#

jax.extend.ffi.register_ffi_target(name, fn, platform='cpu', api_version=1, **kwargs)[源代码][源代码]#

注册一个外部函数目标。

参数:

  • name (str) – 目标的名称。

  • fn (Any) – 一个包含函数指针的 PyCapsule 对象，或一个 dict ，其中键是 FFI 阶段名称（例如 “execute”），值是包含该阶段处理程序指针的 PyCapsule 对象。

  • platform (str) – 目标平台。

  • api_version (int) – 要使用的 XLA 自定义调用 API 版本。支持的版本有：1（默认）用于类型化的 FFI 或 0 用于早期的“自定义调用”API。

  • kwargs (Any) – 任何额外的关键字参数都会直接传递给 register_custom_call_target() 以用于更高级的使用场景。

返回类型:

None

目录