Register a type for serialization. Equivalent to `flax.serialization.from_state_dict`. Args: ty: The type to be registered. ty_to_state_dict: A function that takes an instance of `ty` and returns its state as a dictionary. ty_from_state_dict: A function
(
ty: type, ty_to_state_dict: Callable, ty_from_state_dict: Callable, override: bool = False
)
| 106 | |
| 107 | |
| 108 | def register_serialization_state( |
| 109 | ty: type, ty_to_state_dict: Callable, ty_from_state_dict: Callable, override: bool = False |
| 110 | ): |
| 111 | """Register a type for serialization. |
| 112 | |
| 113 | Equivalent to `flax.serialization.from_state_dict`. |
| 114 | |
| 115 | Args: |
| 116 | ty: The type to be registered. |
| 117 | ty_to_state_dict: A function that takes an instance of `ty` and returns its state as a |
| 118 | dictionary. |
| 119 | ty_from_state_dict: A function that takes an instance of `ty` and a state dict, and returns |
| 120 | a copy of the instance with the restored state. |
| 121 | override: Whether to override a previously registered serialization handler. |
| 122 | """ |
| 123 | if ty in _STATE_DICT_REGISTRY and not override: |
| 124 | raise ValueError(f'A serialization handler for "{ty.__name__}" is already registered.') |
| 125 | _STATE_DICT_REGISTRY[ty] = (ty_to_state_dict, ty_from_state_dict) |
| 126 | |
| 127 | |
| 128 | # Below are serialization implementations for standard container types. |