jax.ffi.register_ffi_type_id

jax.ffi.register_ffi_type_id#

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

Registers a custom type ID for a FFI target.

Parameters:

name (str) – the name of the type ID. This name must be unique within the process.
obj (Any) – a PyCapsule object encapsulating a pointer to the type ID.
platform (str) – the target platform.

Return type:

None