neuralqx.utils.serialization module

A module for serializing/deserializing nearly any Python object, including: - JAX arrays - NumPy arrays - Built-in containers (list, dict, set, tuple) - Arbitrary user-defined classes - Fallback for objects lacking __dict__

It also includes: - A versioning system - Security checks via SHA-256 checksums - Post-load hooks - Optional custom serialization logic with __serializable_attributes__ methods

class SerializerEngine

Bases: object

SerializerEngine is an OOP abstraction that encapsulates all logic for serializing and deserializing objects (including JAX, NumPy, etc.), as well as saving/loading snapshots to/from files with checksums.

The class provides static methods:

• serialize(obj): returns a nested dictionary (pytree)

• deserialize(snapshot): rebuilds the original object from that pytree

• save_to_file(obj, filename): writes serialized data + checksum to disk

• load_from_file(filename): reads data from disk and deserializes it

Internally, it uses helper methods (_serialize_any, _deserialize_any, etc.) to handle specialized types like JAX arrays or NumPy arrays.

static dynamic_import_class(class_name)

Dynamically import a class given its fully qualified string (e.g. “my_package.submodule.MyClass”).

Parameters:

class_name (str) – the dotted path to the class

Return type:

type

Returns:

the actual class object

static compute_sha256(data)

Compute an SHA-256 hash of the given bytes and return it as a hex digest

Parameters:

data (bytes) – the raw bytes to be hashed

Return type:

str

Returns:

hexadecimal SHA-256 digest

static get_fully_qualified_name(obj)

Returns the fully qualified Python name of obj’s class

Parameters:

obj (Any) – the object whose class name we want

Return type:

str

Returns:

e.g. “my_package.submodule.MyClass”

static maybe_post_load_hook(obj)

If the object defines a ‘_post_load_hook()’ method, call it.

This is useful for objects that need to re-initialize ephemeral state.

static get_serializable_state(obj)

Get a dictionary of attributes that should be serialized for this object.

1. If obj.__serializable_attributes__() exists and is callable, use that

2. Otherwise, if obj.__dict__ exists, store all non-callable attributes

3. Otherwise, return an empty dict (meaning no attributes are stored)

Parameters:

obj (Any) – the object to introspect

Return type:

Dict[str, Any]

Returns:

a dictionary of {attr_name: attr_value}

static serialize(obj)

Serialize a Python object into a nested dict (pytree).

The result includes:

• __classname__: the fully qualified class name

• __version__: the module-level serialization format version

• __class_version__: optional version from obj.__class_version__

• __attributes__: the recursively serialized attributes (dict)

:param obj:the object to serialize

Return type:

Dict[str, Any]

Returns:

a nested dictionary with the object’s data

static deserialize(snapshot)

Reconstruct an object from a top-level snapshot (dictionary) previously produced by serialize(…).

:param snapshot:the data structure containing class info and serialized attributes

Return type:

Any

Returns:

the reconstructed (deserialized) object

static save_to_file(obj, filename)

Serialize the object and write it to ‘filename’ using msgpack, along with a SHA-256 checksum to detect tampering. You can also pass it an already serialised dict to directly save

Return type:

None

Parameters:

obj (Union[Any, dict]) – the Python object to be serialized or a serialised dict

:param filename:the path to the file where data will be saved

static load_from_file(filename, raw=False)

Read a snapshot (and its checksum) from ‘filename’ and deserialize it.

Parameters:

• filename (str) – the path to the file containing the serialized data

• raw (bool) – a flag to determine whether the given snapshot is already a serialised dict that we do not want to set into an object

Return type:

Any

Returns:

the reconstructed object from the file

serialize(obj)

Public function to serialize a Python object into a nested dictionary. This delegates to the SerializerEngine class.

Parameters:

obj (Any) – the object to serialize

Return type:

Dict[str, Any]

Returns:

a nested dictionary with the object’s data

deserialize(snapshot)

Public function to reconstruct an object from a snapshot dictionary. Delegates to the SerializerEngine class.

Parameters:

snapshot (Dict[str, Any]) – the data structure containing class info and serialized attributes

Return type:

Any

Returns:

the reconstructed (deserialized) object

save_to_file(obj, filename)

Public function to serialize an object and save it to a file with a checksum. Delegates to the SerializerEngine class.

Parameters:

• obj (Any) – the Python object to serialize

• filename (str) – path to the output file

Return type:

None

load_from_file(filename, raw=False)

Public function to load a snapshot from a file, verify its checksum, and deserialize the object. Delegates to the SerializerEngine class.

Parameters:

• filename (str) – path to the file to read

• raw (bool) – a flag to determine whether the given snapshot is already a serialised dict that we do not want to set into an object

Return type:

Any

Returns:

the reconstructed object

auto_serializable(cls)

A decorator that automatically adds a default __serializable_attributes__() method to a class if none is defined. It will basically return all class attributes.

partial_serialize(obj, attr_names)

Collects a subset of obj’s attributes (those in attr_names) and serializes them into a dictionary that can be saved to disk.

Parameters:

• obj (Any) – the Python object from which we want to extract certain attributes

• attr_names (List[str]) – list of attribute names (strings) to be serialized

Return type:

dict

Returns:

a dict mapping each attribute name to its serialized representation

partial_deserialize(partial_data, existing_obj)

Given a dictionary of serialized attributes (produced by partial_serialize), deserialize them and set them onto existing_obj in place.

Parameters:

• partial_data (dict) – a dict of {attr_name -> serialized_value}

• existing_obj (Any) – the object whose attributes should be updated in-place

Return type:

None

save_multi_partial(obj_map, filename)

For each key in obj_map, we partial-serialize its (obj, attr_list), then store them all in one structure. We’ll also compute a checksum.

Parameters:

• obj_map (Dict[str, Tuple[Any, List[str]]]) –

  a dict like: {

  ”sampler”: (sampler_obj, [“chains”, “rng_state”]), “optimizer”: (optimizer_obj, [“learning_rate”, “momentum”]), …

  }

• filename (str) – The path where we’ll write the msgpack data.

Return type:

None

load_multi_partial(obj_map, filename)

Loads the multi-partial snapshot from ‘filename’. Then for each label in ‘obj_map’, merges the corresponding partial data into the provided object in-place.

Parameters:

• obj_map (Dict[str, Any]) –

  a dict like: {

  ”sampler”: sampler_obj, “optimizer”: optimizer_obj, “network”: network_obj,

  } (Note: no attribute lists here; we already have the partial_data stored in the file)

• filename (str) – path to the msgpack file that has multi_partial_data

Return type:

Dict[str, dict]

Returns:

a dict of {label -> partial_data}, in case we want to do custom merges or just to see the raw data. We do the in-place assignment too