jax.ffi.register_ffi_type

jax.ffi.register_ffi_type(name, type_registration, platform='cpu')

Registers a custom type for a FFI target.

Parameters:

• name (str) – the name of the type. This name must be unique within the process.

• type_registration (TypeRegistration) – a TypeRegistration defining the external type.

• platform (str) – the target platform.

Return type:

None