statedict2pytree 0.2.0__tar.gz → 0.3.0__tar.gz

{statedict2pytree-0.2.0 → statedict2pytree-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: statedict2pytree
-Version: 0.2.0
+Version: 0.3.0
 Summary: Converts torch models into PyTrees for Equinox
 Author-email: "Artur A. Galstyan" <mail@arturgalstyan.dev>
 Requires-Python: ~=3.10
@@ -32,6 +32,12 @@ Usually, if you _declared the fields in the same order as in the PyTorch model_,
 
 (Theoretically, you can rearrange the model in any way you like - e.g. last layer as the first layer - as long as the shapes match!)
 
+## Shape Matching? What's that?
+
+Currently, there is no sophisticated shape matching in place. Two matrices are considered "matching" if the product of their shape match. For example:
+
+1. (8, 1, 1) and (8, ) match, because (8 _ 1 _ 1 = 8)
+
 ## Get Started
 
 ### Installation
@@ -124,3 +130,18 @@ def convert(
 ```
 
 If your models already have the right "order", then you might as well use this function directly. Note that the lists `jax_fields` and `torch_fields` must have the same length and each matching entry must have the same shape!
+
+For the full, automatic experience, use `autoconvert`:
+
+```python
+import statedict2pytree as s2p
+
+my_model = Model(...)
+state_dict = ...
+
+model, state = s2p.autoconvert(my_model, state_dict)
+
+```
+
+This will however only work if your PyTree fields have been declared
+in the same order as they appear in the state dict!

{statedict2pytree-0.2.0 → statedict2pytree-0.3.0}/README.md

@@ -8,6 +8,12 @@ Usually, if you _declared the fields in the same order as in the PyTorch model_,
 
 (Theoretically, you can rearrange the model in any way you like - e.g. last layer as the first layer - as long as the shapes match!)
 
+## Shape Matching? What's that?
+
+Currently, there is no sophisticated shape matching in place. Two matrices are considered "matching" if the product of their shape match. For example:
+
+1. (8, 1, 1) and (8, ) match, because (8 _ 1 _ 1 = 8)
+
 ## Get Started
 
 ### Installation
@@ -100,3 +106,18 @@ def convert(
 ```
 
 If your models already have the right "order", then you might as well use this function directly. Note that the lists `jax_fields` and `torch_fields` must have the same length and each matching entry must have the same shape!
+
+For the full, automatic experience, use `autoconvert`:
+
+```python
+import statedict2pytree as s2p
+
+my_model = Model(...)
+state_dict = ...
+
+model, state = s2p.autoconvert(my_model, state_dict)
+
+```
+
+This will however only work if your PyTree fields have been declared
+in the same order as they appear in the state dict!

{statedict2pytree-0.2.0 → statedict2pytree-0.3.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "statedict2pytree"
-version = "0.2.0"
+version = "0.3.0"
 description = "Converts torch models into PyTrees for Equinox"
 readme = "README.md"
 requires-python = "~=3.10"
@@ -16,7 +16,7 @@ dependencies = [
     "torch",
     "flask",
     "pydantic",
-    "penzai"
+    "penzai",
 ]
 [project.optional-dependencies]
 dev = ["nox", "pre-commit", "pytest", "mkdocs"]

statedict2pytree-0.3.0/pyrightconfig.json

@@ -0,0 +1,4 @@
+{
+  "venvPath": ".",
+  "venv": ".venv"
+}

statedict2pytree-0.3.0/statedict2pytree/__init__.py

@@ -0,0 +1,7 @@
+from statedict2pytree.statedict2pytree import (
+    autoconvert as autoconvert,
+    convert as convert,
+    pytree_to_fields as pytree_to_fields,
+    start_conversion as start_conversion,
+    state_dict_to_fields as state_dict_to_fields,
+)

{statedict2pytree-0.2.0 → statedict2pytree-0.3.0}/statedict2pytree/statedict2pytree.py

@@ -4,6 +4,7 @@ import re
 import equinox as eqx
 import flask
 import jax
+import numpy as np
 from beartype.typing import Optional
 from jaxtyping import PyTree
 from loguru import logger
@@ -31,6 +32,13 @@ PYTREE: Optional[PyTree] = None
 STATE_DICT: Optional[dict] = None
 
 
+def can_reshape(shape1, shape2):
+    product1 = np.prod(shape1)
+    product2 = np.prod(shape2)
+
+    return product1 == product2
+
+
 def get_node(
     tree: PyTree, targets: list[str], log_when_not_found: bool = False
 ) -> PyTree | None:
@@ -146,17 +154,23 @@ def index():
     )
 
 
+def autoconvert(pytree: PyTree, state_dict: dict) -> tuple[PyTree, eqx.nn.State]:
+    jax_fields = pytree_to_fields(pytree)
+    torch_fields = state_dict_to_fields(state_dict)
+    return convert(jax_fields, torch_fields, pytree, state_dict)
+
+
 def convert(
     jax_fields: list[JaxField],
     torch_fields: list[TorchField],
     pytree: PyTree,
     state_dict: dict,
-):
+) -> tuple[PyTree, eqx.nn.State]:
     identity = lambda *args, **kwargs: pytree
     model, state = eqx.nn.make_with_state(identity)()
     state_paths: list[tuple[JaxField, TorchField]] = []
     for jax_field, torch_field in zip(jax_fields, torch_fields):
-        if jax_field.shape != torch_field.shape:
+        if not can_reshape(jax_field.shape, torch_field.shape):
             raise ValueError(
                 "Fields have incompatible shapes!"
                 f"{jax_field.shape=} != {torch_field.shape=}"
@@ -171,7 +185,7 @@ def convert(
                 model = eqx.tree_at(
                     where,
                     model,
-                    state_dict[torch_field.path],
+                    state_dict[torch_field.path].reshape(jax_field.shape),
                 )
     result: dict[str, list[TorchField]] = {}
     for tuple_item in state_paths:

statedict2pytree-0.3.0/tests/test_conv.py

@@ -0,0 +1,61 @@
+import equinox as eqx
+import jax
+import numpy as np
+import statedict2pytree as s2p
+import torch
+
+
+def test_conv():
+    in_channels = 8
+    out_channels = 8
+    kernel_size = 4
+    stride = 2
+    padding = 1
+
+    class J(eqx.Module):
+        conv: eqx.nn.Conv2d
+
+        def __init__(self):
+            self.conv = eqx.nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=padding,
+                key=jax.random.PRNGKey(22),
+            )
+
+    class T(torch.nn.Module):
+        def __init__(self) -> None:
+            super(T, self).__init__()
+            self.conv = torch.nn.Conv2d(
+                in_channels=in_channels,
+                out_channels=out_channels,
+                kernel_size=kernel_size,
+                stride=stride,
+                padding=padding,
+            )
+
+    jax_model = J()
+    torch_model = T()
+    state_dict = torch_model.state_dict()
+
+    jax_fields = s2p.pytree_to_fields(jax_model)
+    torch_fields = s2p.state_dict_to_fields(state_dict)
+
+    model, state = s2p.convert(
+        jax_fields, torch_fields, pytree=jax_model, state_dict=state_dict
+    )
+
+    assert np.allclose(
+        np.array(model.conv.weight), torch_model.conv.weight.detach().numpy()
+    )
+    if torch_model.conv.bias is not None:
+        assert np.allclose(
+            np.array(model.conv.bias),
+            torch_model.conv.bias.detach().numpy().reshape(model.conv.bias.shape),
+        )
+
+
+if __name__ == "__main__":
+    test_conv()

statedict2pytree-0.3.0/tests/test_linear.py

@@ -0,0 +1,45 @@
+import equinox as eqx
+import jax
+import numpy as np
+import statedict2pytree as s2p
+import torch
+
+
+def test_linear():
+    in_features = 10
+    out_features = 10
+
+    class J(eqx.Module):
+        linear: eqx.nn.Linear
+
+        def __init__(self):
+            self.linear = eqx.nn.Linear(
+                in_features, out_features, key=jax.random.PRNGKey(30)
+            )
+
+    class T(torch.nn.Module):
+        def __init__(self) -> None:
+            super(T, self).__init__()
+            self.linear = torch.nn.Linear(in_features, out_features)
+
+    jax_model = J()
+    torch_model = T()
+    state_dict = torch_model.state_dict()
+
+    jax_fields = s2p.pytree_to_fields(jax_model)
+    torch_fields = s2p.state_dict_to_fields(state_dict)
+
+    model, state = s2p.convert(
+        jax_fields, torch_fields, pytree=jax_model, state_dict=state_dict
+    )
+
+    assert np.allclose(
+        np.array(model.linear.weight), torch_model.linear.weight.detach().numpy()
+    )
+    assert np.allclose(
+        np.array(model.linear.bias), torch_model.linear.bias.detach().numpy()
+    )
+
+
+if __name__ == "__main__":
+    test_linear()

statedict2pytree-0.2.0/statedict2pytree/__init__.py

@@ -1,4 +0,0 @@
-from statedict2pytree.statedict2pytree import (
-    convert as convert,
-    start_conversion as start_conversion,
-)