jax.ffi.register_ffi_type_id#

jax.ffi.register_ffi_type_id(name, obj, platform='cpu')[source]#

为 FFI 目标注册自定义类型 ID。

参数:

  • name (str) – 类型 ID 的名称。此名称在进程中必须唯一。

  • obj (Any) – 一个封装指向类型 ID 指针的 PyCapsule 对象。

  • platform (str) – 目标平台。

返回类型:

None