gensbi.utils

Utility functions for GenSBI.

This module provides general utility functions including mathematical operations, model wrapping utilities, plotting functions, and model serialization.

Submodules

• gensbi.utils.math
• gensbi.utils.misc
• gensbi.utils.model_wrapping
• gensbi.utils.plotting
• gensbi.utils.serialization

Functions

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.

Package Contents

gensbi.utils.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.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