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

目录