statedict2pytree 1.0.1__tar.gz → 1.0.2__tar.gz

{statedict2pytree-1.0.1 → statedict2pytree-1.0.2}/PKG-INFO

@@ -1,11 +1,11 @@
-Metadata-Version: 2.3
+Metadata-Version: 2.4
 Name: statedict2pytree
-Version: 1.0.1
+Version: 1.0.2
 Summary: Converts torch models into PyTrees for Equinox
 Author-email: "Artur A. Galstyan" <mail@arturgalstyan.dev>
 Requires-Python: >=3.11
 Requires-Dist: beartype
-Requires-Dist: equinox
+Requires-Dist: equinox>=0.13.0
 Requires-Dist: jax
 Requires-Dist: jaxlib
 Requires-Dist: jaxtyping

statedict2pytree-1.0.2/flax_example.py

@@ -0,0 +1,33 @@
+import torch
+from flax import nnx
+from statedict2pytree import autoconvert, pytree_to_fields
+
+
+class TorchModel(torch.nn.Module):
+    def __init__(self, din, dout):
+        super(TorchModel, self).__init__()
+        self.linear = torch.nn.Linear(din, dout, bias=False)
+        self.linear.weight.data = torch.ones_like(self.linear.weight)
+
+    def forward(self, x):
+        return self.linear(x)
+
+
+class Model(nnx.Module):
+    def __init__(self, din, dout, rngs: nnx.Rngs):
+        self.linear = nnx.Linear(din, dout, rngs=rngs, use_bias=False)
+
+    def __call__(self, x):
+        return self.linear(x)
+
+
+flax_model = Model(2, 64, rngs=nnx.Rngs(0))
+torch_model = TorchModel(2, 64)
+
+pt_fields = pytree_to_fields(flax_model)
+print(pt_fields)
+
+flax_model = autoconvert(flax_model, torch_model.state_dict())
+
+print(flax_model)
+print(flax_model.linear.kernel.value)

{statedict2pytree-1.0.1 → statedict2pytree-1.0.2}/pyproject.toml

@@ -1,13 +1,13 @@
 [project]
 name = "statedict2pytree"
-version = "1.0.1"
+version = "1.0.2"
 description = "Converts torch models into PyTrees for Equinox"
 readme = "README.md"
 requires-python = ">=3.11"
 authors = [{ name = "Artur A. Galstyan", email = "mail@arturgalstyan.dev" }]
 dependencies = [
     "jax",
-    "equinox",
+    "equinox>=0.13.0",
     "jaxlib",
     "beartype",
     "jaxtyping",

{statedict2pytree-1.0.1 → statedict2pytree-1.0.2}/statedict2pytree/converter.py

@@ -94,7 +94,7 @@ def _get_node(
         if len(path) == 0:
             return tree, state_indices
     f, *_ = path
-    if hasattr(tree, "is_stateful"):
+    if hasattr(tree, "is_stateful") and tree.is_stateful():
         if state_indices is None:
             state_indices = {}
         indices = _get_stateindex_fields(tree)
@@ -109,7 +109,10 @@ def _get_node(
     elif isinstance(f, FlattenedIndexKey):
         if isinstance(tree, eqx.nn.State):
             assert state_indices is not None
-            index = state_indices[f.key]
+
+            markers = list(tree._state.keys())
+            marker = markers[f.key]
+            index = state_indices[marker]
             subtree = tree.get(index)
         else:
             subtree = None