Serialisation

feedbax.save (path: str | Path,tree: PyTree[eqx.Module],hyperparameters: Optional[dict] = None) -> None

Save a PyTree to disk along with hyperparameters used to generate it.

Assumes none of the hyperparameters are JAX arrays, as these are not JSON serialisable.

Based on the Equinox serialisation example.

Parameters:

Name	Type	Description	Default
path	str | Path	The path of the file to be saved. Note that the file at this path will be overwritten if it exists.	required
tree	PyTree[Module]	The PyTree to save. Its structure should match the return type of a function setup_func, which will be passed to load.	required
hyperparameters	Optional[dict]	A dictionary of arguments for setup_func that were used to generate the PyTree, and upon loading, will be used to regenerate an appropriate skeleton to populate with the saved values from tree.	None

feedbax.load (path: Path | str,setup_func: Callable[..., PyTree[Any, T]],**kwargs) -> PyTree[Any, T]

Setup a PyTree from stored data and hyperparameters.

Parameters:

Name	Type	Description	Default
path	Path | str	The path of the file to be loaded.	required
setup_func	Callable[..., PyTree[Any, T]]	A function that returns a PyTree of the same structure as the PyTree that was saved to path, and which may take as arguments hyperparameters which save may have saved to the same file. It must take a keyword argument key.	required