safejax.core.save
serialize(params, metadata=None, include_objax_variables=False, filename=None, fs=None)
Serialize JAX, Flax, Haiku, or Objax model params from either FrozenDict
, Dict
, or VarCollection
.
If filename
is not provided, the serialized model is returned as a bytes
object,
otherwise the model is saved to the provided filename
and the filename
is returned.
Parameters:
Name | Type | Description | Default |
---|---|---|---|
params |
ParamsDictLike
|
A |
required |
metadata |
Union[None, Dict[str, str]]
|
A |
None
|
include_objax_variables |
bool
|
Whether to include |
False
|
filename |
Union[PathLike, None]
|
The path to the file where the model params will be saved. |
None
|
fs |
Union[AbstractFileSystem, None]
|
The filesystem to use to save the model params. Defaults to |
None
|
Returns:
Type | Description |
---|---|
Union[bytes, PathLike]
|
The serialized model params as a |
Source code in safejax/core/save.py
Last update:
2023-01-19
Created: 2023-01-19
Created: 2023-01-19