More general pack/unpack for serialization helper

brentyi · Mar 11, 2021 · ba54b59 · ba54b59
1 parent eab5330
commit ba54b59
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/jaxlie/_utils.py b/jaxlie/_utils.py
@@ -1,3 +1,4 @@
+import dataclasses
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -95,10 +96,18 @@ def _unflatten_group(treedef: Any, children: Sequence[Any]) -> T:
 
         # Make object flax-serializable
         def _ty_to_state_dict(x: "MatrixLieGroup") -> Dict[str, types.Array]:
-            return {"params": x.parameters}
+            return {
+                key: flax.serialization.to_state_dict(value)
+                for key, value in vars(x).items()
+            }
 
         def _ty_from_state_dict(x: "MatrixLieGroup", state: Dict) -> "MatrixLieGroup":
-            return type(x)(state["params"])
+            updates: Dict[str, Any] = {}
+            for key, value in vars(x).items():
+                updates[key] = flax.serialization.from_state_dict(
+                    getattr(x, key), value
+                )
+            return dataclasses.replace(x, **updates)
 
         flax.serialization.register_serialization_state(
             ty=cls,