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

statedict2pytree/statedict2pytree.py

@@ -95,7 +95,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 +130,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 +146,21 @@ 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 convert(
+    jax_fields: list[JaxField],
+    torch_fields: list[TorchField],
+    pytree: PyTree,
+    state_dict: dict,
+):
+    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:
+            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 +171,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],
                 )
     result: dict[str, list[TorchField]] = {}
     for tuple_item in state_paths:
@@ -177,7 +184,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.2.0.dist-info/METADATA

@@ -0,0 +1,126 @@
+Metadata-Version: 2.3
+Name: statedict2pytree
+Version: 0.2.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!)
+
+## 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!

{statedict2pytree-0.1.2.dist-info → statedict2pytree-0.2.0.dist-info}/RECORD

@@ -1,8 +1,8 @@
 statedict2pytree/__init__.py,sha256=kMuooLMZQ68rfJSJNVEpJORGnSJFY1sv6jgK9Guh4LY,116
-statedict2pytree/statedict2pytree.py,sha256=LUM19UvNn8R9jau3iYmDLbgfOlznytNEL3J-d-RoSZ0,6455
+statedict2pytree/statedict2pytree.py,sha256=SQ1Xs4VPG_gnjYdAOBzdoNih1RGVsW9x08UQGkIFUdg,6640
 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,,
+statedict2pytree-0.2.0.dist-info/METADATA,sha256=x9GhcbG0io5HIpfwocS-nsPYyKRTmG90uLKyckSm394,3673
+statedict2pytree-0.2.0.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
+statedict2pytree-0.2.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
-
-```