eformer.pytree._serialization

eformer.pytree._serialization.compress_bytes(data: bytes) → bytes

Compress serialized bytes using gzip.

Parameters

data – Bytes to compress

Returns

Compressed bytes

eformer.pytree._serialization.current_path()

Get the current state_dict path during deserialization.

Returns

Current path as a slash-separated string, useful for error messages.

Return type

str

Examples

>>>
>>> path = current_path()
>>>

eformer.pytree._serialization.decompress_bytes(data: bytes) → bytes

Decompress gzipped bytes.

Parameters

data – Compressed bytes

Returns

Decompressed bytes

eformer.pytree._serialization.from_bytes(target, encoded_bytes: bytes)

Restore optimizer or other object from msgpack-serialized state-dict.

Parameters

• target – template object with state-dict registrations that matches the structure being deserialized from encoded_bytes.

• encoded_bytes – msgpack serialized object structurally isomorphic to target. Typically a flax model or optimizer.

Returns

A new object structurally isomorphic to target containing the updated leaf data from saved data.

eformer.pytree._serialization.from_compressed_bytes(target, compressed_data: bytes)

Load pytree from compressed msgpack bytes.

Parameters

• target – Template pytree

• compressed_data – Compressed msgpack bytes

Returns

Loaded pytree

eformer.pytree._serialization.from_state_dict(target, state: dict[str, Any], name: str = '.')

Restores the state of the given target using a state dict.

This function takes the current target as an argument. This lets us know the exact structure of the target, as well as lets us add assertions that shapes and dtypes don’t change.

In practice, none of the leaf values in target are actually used. Only the tree structure, shapes and dtypes.

Parameters

• target – the object of which the state should be restored.

• state – a dictionary generated by to_state_dict with the desired new state for target.

• name – name of branch taken, used to improve deserialization error messages.

Returns

A copy of the object with the restored state.

eformer.pytree._serialization.get_serialization_info(pytree) → dict[str, Any]

Get information about pytree serialization.

Parameters

pytree – PyTree to analyze

Returns

Dictionary with serialization information

eformer.pytree._serialization.is_serializable(target)

Check if a target object is serializable.

Parameters

target – Object or type to check for serializability.

Returns

True if the target type is registered for serialization.

Return type

bool

Examples

>>> is_serializable(my_model)
True
>>> is_serializable(42)
False

eformer.pytree._serialization.load_from_file(target, filepath: str)

Load pytree from a file.

Parameters

• target – Template pytree for structure

• filepath – Path to load from

Returns

Loaded pytree with same structure as target

eformer.pytree._serialization.msgpack_restore(encoded_pytree: bytes)

Restore data structure from bytes in msgpack format.

Low-level function that only supports python trees with array leaves, for custom objects use from_bytes.

Parameters

encoded_pytree – msgpack-encoded bytes of python tree.

Returns

Python tree of dict, list, tuple with python primitive and array leaves.

eformer.pytree._serialization.msgpack_serialize(pytree, in_place: bool = False) → bytes

Save data structure to bytes in msgpack format.

Low-level function that only supports python trees with array leaves, for custom objects use to_bytes. It splits arrays above MAX_CHUNK_SIZE into multiple chunks.

Parameters

• pytree – python tree of dict, list, tuple with python primitives and array leaves.

• in_place – boolean specifying if pytree should be modified in place.

Returns

msgpack-encoded bytes of pytree.

eformer.pytree._serialization.register_serialization_state(ty, ty_to_state_dict, ty_from_state_dict, override=False)

Register a type for serialization.

Parameters

• 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 that takes an instance of ty and a state dict, and returns a copy of the instance with the restored state.

• override – override a previously registered serialization handler (default: False).

eformer.pytree._serialization.save_to_file(target, filepath: str) → None

Save pytree to a file.

Parameters

• target – PyTree to save

• filepath – Path to save file

eformer.pytree._serialization.to_bytes(target) → bytes

Save optimizer or other object as msgpack-serialized state-dict.

Parameters

target – template object with state-dict registrations to be serialized to msgpack format. Typically a flax model or optimizer.

Returns

Bytes of msgpack-encoded state-dict of target object.

eformer.pytree._serialization.to_compressed_bytes(target) → bytes

Save pytree as compressed msgpack bytes.

Parameters

target – PyTree to serialize

Returns

Compressed msgpack bytes

eformer.pytree._serialization.to_state_dict(target) → dict[str, Any]

Convert an object to a state dictionary.

Parameters

target – Object to convert, typically a PyTree with registered types.

Returns

State dictionary representation of the target.

Return type

dict

Note

Only registered types are converted; others are returned as-is.

Examples

>>> model = MyModel()
>>> state = to_state_dict(model)
>>>

eformer.pytree._serialization.validate_serializable(pytree) → tuple[bool, list[str]]

Validate if a pytree is fully serializable.

Parameters

pytree – PyTree to validate

Returns

Tuple of (is_valid, list_of_issues)