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