statedict2pytree 0.2.0__py3-none-any.whl → 0.4.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:
@@ -38,7 +46,7 @@ def get_node(
         return tree
     else:
         next_target: str = targets[0]
-        if bool(re.search(r"\[\d\]", next_target)):
+        if bool(re.search(r"\[\d+\]", next_target)):
             split_index = next_target.rfind("[")
             name, index = next_target[:split_index], next_target[split_index:]
             index = index[1:-1]
@@ -146,17 +154,26 @@ 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)
+
+    for k, v in state_dict.items():
+        state_dict[k] = v.numpy()
+    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 +188,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.2.0.dist-info → statedict2pytree-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: statedict2pytree
-Version: 0.2.0
+Version: 0.4.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.4.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+statedict2pytree/__init__.py,sha256=lXxSaFFvkhXweXp5oHSkg_dPjdp49OsF8xoqwX4d_4E,240
+statedict2pytree/statedict2pytree.py,sha256=u1PddUBY_MHErl3tstSpJt2a6a_H-RizvxP9anPoFpQ,7175
+statedict2pytree/static/input.css,sha256=zBp60NAZ3bHTLQ7LWIugrCbOQdhiXdbDZjSLJfg6KOw,59
+statedict2pytree/static/output.css,sha256=KZ9GzeV3q0XKjbEiTdPkC6yV-R6jzXRflRm2S16VkJA,40813
+statedict2pytree/templates/index.html,sha256=0uG3dB2pAa1f2wcfTpYSO7TBNL77i2ALJP5rIhsbEnk,7506
+statedict2pytree-0.4.0.dist-info/METADATA,sha256=BbJaUhVv-Qb4-e2DlX0WBNgPOKsTjlKv8ISJZgFNx1o,4232
+statedict2pytree-0.4.0.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
+statedict2pytree-0.4.0.dist-info/RECORD,,

statedict2pytree-0.2.0.dist-info/RECORD

@@ -1,8 +0,0 @@
-statedict2pytree/__init__.py,sha256=kMuooLMZQ68rfJSJNVEpJORGnSJFY1sv6jgK9Guh4LY,116
-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.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,,