statedict2pytree 0.1.2__py3-none-any.whl → 0.3.0__py3-none-any.whl

statedict2pytree/__init__.py

@@ -1,4 +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/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:
@@ -95,7 +103,7 @@ def visualize_with_penzai():
         return flask.jsonify({"error": "No data received"})
     jax_fields = request_data["jaxFields"]
     torch_fields = request_data["torchFields"]
-    model, state = convert(jax_fields, torch_fields)
+    model, state = convert(jax_fields, torch_fields, PYTREE, STATE_DICT)
     with pz.ts.active_autovisualizer.set_scoped(pz.ts.ArrayAutovisualizer()):
         html_jax = pz.ts.render_to_html((model, state))
         html_torch = pz.ts.render_to_html(STATE_DICT)
@@ -130,7 +138,7 @@ def convert_torch_to_jax():
         torch_fields.append(TorchField(path=f["path"], shape=shape_tuple))
 
     name = request_data["name"]
-    model, state = convert(jax_fields, torch_fields)
+    model, state = convert(jax_fields, torch_fields, PYTREE, STATE_DICT)
     eqx.tree_serialise_leaves(name, (model, state))
 
     return flask.jsonify({"status": "success"})
@@ -146,14 +154,27 @@ def index():
     )
 
 
-def convert(jax_fields: list[JaxField], torch_fields: list[TorchField]):
-    global PYTREE, STATE_DICT
-    if STATE_DICT is None:
-        raise ValueError("STATE_DICT must not be None!")
-    identity = lambda *args, **kwargs: PYTREE
+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 not can_reshape(jax_field.shape, torch_field.shape):
+            raise ValueError(
+                "Fields have incompatible shapes!"
+                f"{jax_field.shape=} != {torch_field.shape=}"
+            )
         path = jax_field.path.split(".")[1:]
         if "StateIndex" in jax_field.type:
             state_paths.append((jax_field, torch_field))
@@ -164,7 +185,7 @@ def convert(jax_fields: list[JaxField], torch_fields: list[TorchField]):
                 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:
@@ -177,7 +198,7 @@ def convert(jax_fields: list[JaxField], torch_fields: list[TorchField]):
     for key in result:
         state_index = get_node(model, key.split("."))
         if state_index is not None:
-            to_replace_tuple = tuple([STATE_DICT[i.path] for i in result[key]])
+            to_replace_tuple = tuple([state_dict[i.path] for i in result[key]])
             state = state.set(state_index, to_replace_tuple)
     return model, state
 

statedict2pytree-0.3.0.dist-info/METADATA

@@ -0,0 +1,147 @@
+Metadata-Version: 2.3
+Name: statedict2pytree
+Version: 0.3.0
+Summary: Converts torch models into PyTrees for Equinox
+Author-email: "Artur A. Galstyan" <mail@arturgalstyan.dev>
+Requires-Python: ~=3.10
+Requires-Dist: beartype
+Requires-Dist: equinox>=0.11.4
+Requires-Dist: flask
+Requires-Dist: jax
+Requires-Dist: jaxlib
+Requires-Dist: jaxtyping
+Requires-Dist: loguru
+Requires-Dist: penzai
+Requires-Dist: pydantic
+Requires-Dist: torch
+Requires-Dist: typing-extensions
+Provides-Extra: dev
+Requires-Dist: mkdocs; extra == 'dev'
+Requires-Dist: nox; extra == 'dev'
+Requires-Dist: pre-commit; extra == 'dev'
+Requires-Dist: pytest; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# statedict2pytree
+
+![statedict2pytree](torch2jax.png "A ResNet demo")
+
+The goal of this package is to simplify the conversion from PyTorch models into JAX PyTrees (which can be used e.g. in Equinox). The way this works is by putting both models side my side and aligning the weights in the right order. Then, all statedict2pytree is doing, is iterating over both lists and matching the weight matrices.
+
+Usually, if you _declared the fields in the same order as in the PyTorch model_, you don't have to rearrange anything -- but the option is there if you need it.
+
+(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
+
+Run
+
+```bash
+pip install statedict2pytree
+
+```
+
+### Basic Example
+
+```python
+import equinox as eqx
+import jax
+import torch
+import statedict2pytree as s2p
+
+
+def test_mlp():
+    in_size = 784
+    out_size = 10
+    width_size = 64
+    depth = 2
+    key = jax.random.PRNGKey(22)
+
+    class EqxMLP(eqx.Module):
+        mlp: eqx.nn.MLP
+        batch_norm: eqx.nn.BatchNorm
+
+        def __init__(self, in_size, out_size, width_size, depth, key):
+            self.mlp = eqx.nn.MLP(in_size, out_size, width_size, depth, key=key)
+            self.batch_norm = eqx.nn.BatchNorm(out_size, axis_name="batch")
+
+        def __call__(self, x, state):
+            return self.batch_norm(self.mlp(x), state)
+
+    jax_model = EqxMLP(in_size, out_size, width_size, depth, key)
+
+    class TorchMLP(torch.nn.Module):
+        def __init__(self, in_size, out_size, width_size, depth):
+            super(TorchMLP, self).__init__()
+            self.layers = torch.nn.ModuleList()
+            self.layers.append(torch.nn.Linear(in_size, width_size))
+            for _ in range(depth - 1):
+                self.layers.append(torch.nn.Linear(width_size, width_size))
+            self.layers.append(torch.nn.Linear(width_size, out_size))
+            self.batch_norm = torch.nn.BatchNorm1d(out_size)
+
+        def forward(self, x):
+            for layer in self.layers[:-1]:
+                x = torch.relu(layer(x))
+            x = self.batch_norm(self.layers[-1](x))
+            return x
+
+    torch_model = TorchMLP(in_size, out_size, width_size, depth)
+    state_dict = torch_model.state_dict()
+    s2p.start_conversion(jax_model, state_dict)
+
+
+if __name__ == "__main__":
+    test_mlp()
+
+```
+
+There exists also a function called `s2p.convert` which does the actual conversion:
+
+```python
+
+class Field(BaseModel):
+    path: str
+    shape: tuple[int, ...]
+
+
+class TorchField(Field):
+    pass
+
+
+class JaxField(Field):
+    type: str
+
+def convert(
+    jax_fields: list[JaxField],
+    torch_fields: list[TorchField],
+    pytree: PyTree,
+    state_dict: dict,
+):
+...
+```
+
+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.3.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+statedict2pytree/__init__.py,sha256=lXxSaFFvkhXweXp5oHSkg_dPjdp49OsF8xoqwX4d_4E,240
+statedict2pytree/statedict2pytree.py,sha256=X5Ljf4lYhhH7_V4KgdciChncbTt7YZpIWHcOxcZ3l48,7103
+statedict2pytree/static/input.css,sha256=zBp60NAZ3bHTLQ7LWIugrCbOQdhiXdbDZjSLJfg6KOw,59
+statedict2pytree/static/output.css,sha256=KZ9GzeV3q0XKjbEiTdPkC6yV-R6jzXRflRm2S16VkJA,40813
+statedict2pytree/templates/index.html,sha256=0uG3dB2pAa1f2wcfTpYSO7TBNL77i2ALJP5rIhsbEnk,7506
+statedict2pytree-0.3.0.dist-info/METADATA,sha256=YSK4tWzNQemyZ1xKq5BhWiLWWc-RDr4E9q_eV_iOsdw,4232
+statedict2pytree-0.3.0.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
+statedict2pytree-0.3.0.dist-info/RECORD,,

statedict2pytree-0.1.2.dist-info/METADATA

@@ -1,43 +0,0 @@
-Metadata-Version: 2.3
-Name: statedict2pytree
-Version: 0.1.2
-Summary: Converts torch models into PyTrees for Equinox
-Author-email: "Artur A. Galstyan" <mail@arturgalstyan.dev>
-Requires-Python: ~=3.10
-Requires-Dist: beartype
-Requires-Dist: equinox>=0.11.4
-Requires-Dist: flask
-Requires-Dist: jax
-Requires-Dist: jaxlib
-Requires-Dist: jaxtyping
-Requires-Dist: loguru
-Requires-Dist: pydantic
-Requires-Dist: torch
-Requires-Dist: typing-extensions
-Provides-Extra: dev
-Requires-Dist: mkdocs; extra == 'dev'
-Requires-Dist: nox; extra == 'dev'
-Requires-Dist: pre-commit; extra == 'dev'
-Requires-Dist: pytest; extra == 'dev'
-Description-Content-Type: text/markdown
-
-# statedict2pytree
-
-![statedict2pytree](torch2jax.png "A ResNet demo")
-
-The goal of this package is to simplify the conversion from PyTorch models into JAX PyTrees (which can be used e.g. in Equinox). The way this works is by putting both models side my side and aligning the weights in the right order. Then, all statedict2pytree is doing, is iterating over both lists and matching the weight matrices.
-
-Usually, if you _declared the fields in the same order as in the PyTorch model_, you don't have to rearrange anything -- but the option is there if you need it.
-
-(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!)
-
-## Get Started
-
-### Installation
-
-Run
-
-```bash
-pip install statedict2pytree
-
-```

statedict2pytree-0.1.2.dist-info/RECORD

@@ -1,8 +0,0 @@
-statedict2pytree/__init__.py,sha256=kMuooLMZQ68rfJSJNVEpJORGnSJFY1sv6jgK9Guh4LY,116
-statedict2pytree/statedict2pytree.py,sha256=LUM19UvNn8R9jau3iYmDLbgfOlznytNEL3J-d-RoSZ0,6455
-statedict2pytree/static/input.css,sha256=zBp60NAZ3bHTLQ7LWIugrCbOQdhiXdbDZjSLJfg6KOw,59
-statedict2pytree/static/output.css,sha256=KZ9GzeV3q0XKjbEiTdPkC6yV-R6jzXRflRm2S16VkJA,40813
-statedict2pytree/templates/index.html,sha256=0uG3dB2pAa1f2wcfTpYSO7TBNL77i2ALJP5rIhsbEnk,7506
-statedict2pytree-0.1.2.dist-info/METADATA,sha256=X-79GNzLPC9VXRPSVTJao9ysWygmiseBKLd4GmgAY-g,1437
-statedict2pytree-0.1.2.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
-statedict2pytree-0.1.2.dist-info/RECORD,,