Add tests/utils.py and tests/__init__.py to import it as import .utils
Implement assert_over_trees function in tests/utils.py
Add chex assertions to jax, flax, and haiku params
🔗 Linked Issue/s
24
📌 Notes
jaxtyping is another potential library to assert over trees
objax.variable.VarCollection is not yet supported as still needs more refinement
🧪 Tests
[X] Did you implement unit tests if required?
If the above checkbox is checked, describe how you unit-tested it.
Add assertions over trees with chex from @deepmind to make sure that the original params dict matches the one serialized and deserialized using safejax with safetensors as the tensor storage format.
✨ Features
tests/utils.pyandtests/__init__.pyto import it asimport .utilsassert_over_treesfunction intests/utils.pychexassertions tojax,flax, andhaikuparams🔗 Linked Issue/s
24
📌 Notes
jaxtypingis another potential library to assert over treesobjax.variable.VarCollectionis not yet supported as still needs more refinement🧪 Tests
If the above checkbox is checked, describe how you unit-tested it.
Add assertions over trees with
chexfrom @deepmind to make sure that the original params dict matches the one serialized and deserialized usingsafejaxwithsafetensorsas the tensor storage format.