gensbi.utils.serialization

Portable safetensors save/load for flax nnx models.

Exports the weights of any flax.nnx.Module to a single, framework-neutral .safetensors file, and loads them back into a model the caller has already reconstructed from its Params. The file stores a flat {str: array} table (nnx state paths joined with ".") plus a small provenance metadata blob; it does not carry enough information to rebuild the model architecture (an explicit non-goal).

Attributes

_DEFAULT_METADATA
_SEP

Functions

_flat_arrays(model, wrt)	Flatten model state to {tuple_path: array} (ints preserved).
_join_key(path)	Join an nnx state-path tuple into a safetensors string key.
load_safetensors(model, path, *[, strict])	Load weights from a .safetensors file into model in place.
save_safetensors(model, path, *[, metadata, wrt])	Save model's weights to a single .safetensors file.

Module Contents

gensbi.utils.serialization._flat_arrays(model, wrt)

Flatten model state to {tuple_path: array} (ints preserved).

Return type:

dict[tuple, Any]

gensbi.utils.serialization._join_key(path)

Join an nnx state-path tuple into a safetensors string key.

Integer nnx.List indices are stringified; a non-integer component that contains the "." separator is unrepresentable and raises ValueError.

Parameters:

path (tuple)

Return type:

str

gensbi.utils.serialization.load_safetensors(model, path, *, strict=True)

Load weights from a .safetensors file into model in place.

The caller must have reconstructed model from its Params first; that model is the structural schema.

Parameters:

• model (nnx.Module) – Target model, rebuilt from its Params.

• path (str | os.PathLike) – Source .safetensors file.

• strict (bool, optional) – If True (default), the file’s key set must equal the model’s and every shared key must match shape (ValueError otherwise). If False, only the intersection is loaded; model leaves absent from the file keep their current values and file keys absent from the model are ignored.

Returns:

model – The same model object, updated in place with the loaded weights.

Return type:

nnx.Module

gensbi.utils.serialization.save_safetensors(model, path, *, metadata=None, wrt=None)

Save model’s weights to a single .safetensors file.

Parameters:

• model (nnx.Module) – Any flax nnx module.

• path (str | os.PathLike) – Destination .safetensors file.

• metadata (mapping, optional) – Extra provenance, stringified and merged over (overriding) the defaults format/version/framework/model_class.

• wrt (nnx filter, optional) – Restrict the saved variable collections (e.g. nnx.Param). Default saves the full state.

Return type:

None

gensbi.utils.serialization._DEFAULT_METADATA

gensbi.utils.serialization._SEP = '.'