jaxlib.xla_client.注册自定义调用目标

jaxlib.xla_client.注册自定义调用目标#

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

目录