Split flax.py into save.py, load.py, and utils.py for readability
save.py contains serialize
load.py contains deserialize
utils.py contains both flatten_dict and unflatten_dict
Add freeze_dict param to unflatten_dict to either convert it to FrozenDict or keep it as a Dict (used for flax)
Update unit tests with pytest to cover every safejax function
Test dm-haiku model param serialization over haiku.nets.ResNet50
Add more examples/ for both flax and dm-haiku
🐛 Bug Fixes
Fix bug while unflattening dictionaries in unflatten_dict due to a variable being overwritten
🧪 Tests
[X] Did you implement unit tests if required?
If the above checkbox is checked, describe how you unit-tested it.
Add some assertions to make sure both safejax.utils.flatten_dict and safejax.utils.unflatten_dict work as expected to avoid bug mentioned above with unflatten_dict
Add some more unit tests for safejax.load and safejax.save due to the recent split of both files
✨ Features
flax.pyintosave.py,load.py, andutils.pyfor readabilitysave.pycontainsserializeload.pycontainsdeserializeutils.pycontains bothflatten_dictandunflatten_dictfreeze_dictparam tounflatten_dictto either convert it toFrozenDictor keep it as aDict(used forflax)pytestto cover everysafejaxfunctiondm-haikumodel param serialization overhaiku.nets.ResNet50examples/for bothflaxanddm-haiku🐛 Bug Fixes
unflatten_dictdue to a variable being overwritten🧪 Tests
If the above checkbox is checked, describe how you unit-tested it.
safejax.utils.flatten_dictandsafejax.utils.unflatten_dictwork as expected to avoid bug mentioned above withunflatten_dictsafejax.loadandsafejax.savedue to the recent split of both files