⚡️ Overall improvements, bug fixes, more unit tests, and `dm-haiku` compatibility tested #6

alvarobartt · 2022-12-26T10:27:44Z

✨ Features

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

Fix bug while unflattening dictionaries in unflatten_dict due to a variable being overwritten

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

As it works with JAX and Haiku too, not just Flax. So we're extending `safejax` use-cases.

Already tested with both `flax` and `haiku` models out of the box, so `safejax` is out of `0.2.0.dev0` :)

alvarobartt added 18 commits December 25, 2022 11:48

🔖 Set 0.2.0.dev0 version

a8ffb49

⚡️ Split flax.py in load.py and save.py

206f650

As it works with JAX and Haiku too, not just Flax. So we're extending `safejax` use-cases.

🧪 Update safejax imports in tests/

249bf38

🩹 Update safejax outdated imports

feae78c

⚡️ Add freeze_dict param in unflatten_dict

2cf1eb0

🐛 Fix bug as params were not properly unflattened

5aa5fa4

♻️ Move flatten_dict and unflatten_dict to utils.py

2912b2a

🔥 Remove supported_value_types param

b1cf39b

⚗️ Update safejax unit tests

f0ceb61

📝 Update README.md

a06244a

💡 Update safejax docstrings

27db1e2

💡 Update safejax description

226119b

📝 Update README.md

a73302a

✨ Update and add more examples/

e195003

📦 Add author.email in pyproject.toml

6239fb4

🎨 Format pyproject.toml using be5invis.toml

ab7b6cf

🔖 Set 0.2.0 version

e2e8f3f

Already tested with both `flax` and `haiku` models out of the box, so `safejax` is out of `0.2.0.dev0` :)

✏️ Fix typos and grammar (thanks @grammarly)

9a0b658

alvarobartt self-assigned this Dec 26, 2022

alvarobartt merged commit aaee170 into main Dec 26, 2022

alvarobartt deleted the extend-usage branch December 26, 2022 10:34