jaxlib.xla_client.register_custom_call_target

目录

jaxlib.xla_client.register_custom_call_target#

jaxlib.xla_client.register_custom_call_target(name, fn, platform='cpu', api_version=0, traits=<CustomCallTargetTraits.DEFAULT: 0>)#

注册自定义调用目标。

参数:

name (str) – 包含函数名称的字节。
fn (Any) – 包含函数指针的 PyCapsule 对象。
platform (str) – 目标平台。
api_version (int) – 要使用的 XLA FFI 版本。支持的版本有：0 用于非类型化 FFI，1 用于类型化 FFI。
traits (CustomCallTargetTraits) – 与 XLA FFI 处理程序特性对应的自定义调用特性。

返回类型:

None