jax.export.register_namedtuple_serialization

jax.export.register_namedtuple_serialization(nodetype, *, serialized_name)

註冊 namedtuple 以進行序列化和反序列化。

JAX 對於 collections.namedtuple 具有原生 PyTree 支援，並且不需要呼叫 jax.tree_util.register_pytree_node。但是，如果您想要序列化具有 namedtuple 型別的輸入或輸出的函數，您必須註冊該型別以進行序列化。

參數:

• nodetype (type[T]) – 我們想要序列化其 PyTree 節點的型別。嘗試為 nodetype 註冊多個序列化是錯誤的。在反序列化時，此型別必須具有與序列化期間存在的相同金鑰集。

• serialized_name (str) – 一個字串，將出現在序列化中，並將用於在反序列化期間查找註冊。嘗試為 serialized_name 註冊多個序列化是錯誤的。

返回:

與作為 nodetype 傳遞的型別相同，以便此函數可以用作類別裝飾器。

返回型別:

type[T]