statedict2pytree 0.6.0__py3-none-any.whl → 1.0.1__py3-none-any.whl

statedict2pytree/__init__.py

@@ -0,0 +1,8 @@
+from .converter import (
+    autoconvert,  # noqa
+    convert,  # noqa
+    is_numerical,  # noqa
+    move_running_fields_to_the_end,  # noqa
+    pytree_to_fields,  # noqa
+    state_dict_to_fields,  # noqa
+)  # noqa

statedict2pytree/converter.py

@@ -0,0 +1,293 @@
+import os
+import pathlib
+import tempfile
+
+import equinox as eqx
+import jax
+import jax.numpy as jnp
+import numpy as np
+from beartype.typing import Any, Callable
+from jax.tree_util import FlattenedIndexKey, GetAttrKey, KeyPath, SequenceKey
+from jaxtyping import Array, PyTree
+from pydantic import BaseModel
+from tqdm import tqdm
+
+
+class ChunkifiedPytreePath(BaseModel):
+    path: str
+
+
+class ChunkifiedStatedictPath(BaseModel):
+    path: str
+
+
+class TorchField(BaseModel):
+    path: str
+    shape: tuple[int, ...]
+    skip: bool = False
+
+
+class JaxField(BaseModel):
+    path: KeyPath
+    shape: tuple[int, ...]
+    skip: bool = False
+
+
+def is_numerical(element: Any):
+    if hasattr(element, "dtype"):
+        # Check if it's a JAX or NumPy array
+        return (
+            np.issubdtype(element.dtype, np.integer)
+            or np.issubdtype(element.dtype, np.floating)
+            or np.issubdtype(element.dtype, np.complexfloating)
+        ) and not np.issubdtype(element.dtype, np.bool_)
+    return False
+
+
+def _default_floating_dtype():
+    if jax.config.jax_enable_x64:  # pyright: ignore
+        return jnp.float64
+    else:
+        return jnp.float32
+
+
+def _can_reshape(shape1: tuple, shape2: tuple):
+    """
+    Check if two shapes can be reshaped to each other.
+
+    Args:
+        shape1 (tuple): First shape.
+        shape2 (tuple): Second shape.
+
+    Returns:
+        bool: True if shapes can be reshaped to each other, False otherwise.
+    """
+    product1 = np.prod(shape1)
+    product2 = np.prod(shape2)
+
+    return product1 == product2
+
+
+def _get_stateindex_fields(obj) -> dict:
+    state_indices = {}
+
+    for attr_name in dir(obj):
+        if attr_name.startswith("_"):
+            continue
+
+        try:
+            attr_value = getattr(obj, attr_name)
+            if isinstance(attr_value, eqx.nn.StateIndex):
+                state_indices[attr_name] = attr_value
+        except:  # noqa
+            pass
+
+    return state_indices
+
+
+def _get_node(
+    tree: PyTree, path: KeyPath, state_indices: dict | None = None
+) -> tuple[PyTree | None, dict | None]:
+    if tree is None:
+        return None, {}
+    else:
+        if len(path) == 0:
+            return tree, state_indices
+    f, *_ = path
+    if hasattr(tree, "is_stateful"):
+        if state_indices is None:
+            state_indices = {}
+        indices = _get_stateindex_fields(tree)
+        for attr_name in indices:
+            index: eqx.nn.StateIndex = indices[attr_name]
+            assert isinstance(index, eqx.nn.StateIndex)
+            state_indices[index.marker] = index
+    if isinstance(f, SequenceKey):
+        subtree = tree[f.idx]
+    elif isinstance(f, GetAttrKey):
+        subtree = getattr(tree, f.name)
+    elif isinstance(f, FlattenedIndexKey):
+        if isinstance(tree, eqx.nn.State):
+            assert state_indices is not None
+            index = state_indices[f.key]
+            subtree = tree.get(index)
+        else:
+            subtree = None
+    else:
+        subtree = None
+    return _get_node(subtree, path[1:], state_indices)
+
+
+def _replace_node(
+    tree: PyTree, path: KeyPath, new_value: Array, state_indices: dict | None = None
+) -> PyTree:
+    def where_wrapper(t):
+        node, _ = _get_node(t, path=path, state_indices=state_indices)
+        return node
+
+    node, _ = _get_node(tree, path=path, state_indices=state_indices)
+
+    if node is not None and eqx.is_array(node):
+        tree = eqx.tree_at(
+            where_wrapper,
+            tree,
+            new_value.reshape(node.shape),
+        )
+    else:
+        print("WARNING: Couldn't find: ", jax.tree_util.keystr(path))
+    return tree
+
+
+def move_running_fields_to_the_end(
+    torchfields: list[TorchField], identifier: str = "running_"
+):
+    """
+    Helper function to move fields that contain the given string to the end of the
+    torchfields. Helpful for stateful layers such as BatchNorm, which appear at the
+    end of Equinox pytrees.
+    """
+    i = 0
+    total = 0
+    while i + total < len(torchfields):
+        if identifier in torchfields[i].path:
+            field = torchfields.pop(i)
+            torchfields.append(field)
+            total += 1
+        else:
+            i += 1
+    return torchfields
+
+
+def state_dict_to_fields(
+    state_dict: dict[str, Any],
+) -> list[TorchField]:
+    if state_dict is None:
+        return []
+    fields: list[TorchField] = []
+    for key, value in state_dict.items():
+        if hasattr(value, "shape") and len(value.shape) > 0:
+            fields.append(TorchField(path=key, shape=tuple(value.shape)))
+    return fields
+
+
+def pytree_to_fields(
+    pytree: PyTree,
+    model_order: list[str] | None = None,
+    filter: Callable[[Array], bool] = eqx.is_array,
+) -> tuple[list[JaxField], dict | None]:
+    jaxfields = []
+    paths = jax.tree.leaves_with_path(pytree)
+    i = {}
+    for p in paths:
+        keys, _ = p
+        n, i = _get_node(pytree, keys, i)
+        if n is not None and filter(n):
+            jaxfields.append(JaxField(path=keys, shape=n.shape))
+
+    if model_order is not None:
+        ordered_jaxfields = []
+        path_dict = {jax.tree_util.keystr(field.path): field for field in jaxfields}
+
+        for path_str in model_order:
+            if path_str in path_dict:
+                ordered_jaxfields.append(path_dict[path_str])
+                del path_dict[path_str]
+        ordered_jaxfields.extend(path_dict.values())
+        jaxfields = ordered_jaxfields
+    return jaxfields, i
+
+
+def _chunkify_state_dict(
+    state_dict: dict[str, np.ndarray], target_path: str
+) -> ChunkifiedStatedictPath:
+    """
+    Convert a PyTorch state dict into chunked files and save them to the specified path.
+
+    Args:
+        state_dict (dict[str, np.ndarray]): The PyTorch state dict to be chunked.
+        target_path (str): The directory where chunked files will be saved.
+
+    Returns:
+        ChunkifiedStatedictPath: A path to the chunked files
+    """
+
+    for key in state_dict.keys():
+        if not hasattr(state_dict[key], "shape"):
+            continue
+        path = pathlib.Path(target_path) / "state_dict"
+
+        if not os.path.exists(path):
+            os.mkdir(path)
+        np.save(path / key, state_dict[key])
+
+    return ChunkifiedStatedictPath(path=str(pathlib.Path(target_path)))
+
+
+def convert(
+    state_dict: dict[str, Any],
+    pytree: PyTree,
+    jaxfields: list[JaxField],
+    state_indices: dict | None,
+    torchfields: list[TorchField],
+    dtype: Any | None = None,
+) -> PyTree:
+    if dtype is None:
+        dtype = _default_floating_dtype()
+    assert dtype is not None
+    state_dict_np: dict[str, np.ndarray] = {
+        k: state_dict[k].detach().numpy() for k in state_dict
+    }
+
+    for k in state_dict_np:
+        if np.issubdtype(state_dict_np[k].dtype, np.floating):
+            state_dict_np[k] = state_dict_np[k].astype(dtype)
+
+    if len(torchfields) != len(jaxfields):
+        raise ValueError(
+            f"Length of state_dict ({len(torchfields)}) "
+            f"!= length of pytree ({len(jaxfields)})"
+        )
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        chunkified_statedict_path = _chunkify_state_dict(state_dict_np, tmpdir)
+        del state_dict_np, state_dict
+        for t, j in tqdm(zip(torchfields, jaxfields), total=len(torchfields)):
+            if not _can_reshape(t.shape, j.shape):
+                raise ValueError(
+                    f"Cannot reshape {t.shape} "
+                    f"into shape {j.shape}. "
+                    "Note that the order of the fields matters "
+                    "and that you can mark arrays as skippable. "
+                    f"{t.path=} "
+                    f"{jax.tree_util.keystr(j.path)=}"
+                )
+            state_dict_dir = pathlib.Path(chunkified_statedict_path.path) / "state_dict"
+            filename = state_dict_dir / t.path
+            new_value = jnp.array(np.load(str(filename) + ".npy"))
+
+            n, _ = _get_node(pytree, j.path, state_indices)
+            assert n is not None, f"Node {j.path} not found"
+            assert _can_reshape(n.shape, new_value.shape), (
+                f"Cannot reshape {n.shape} into {new_value.shape}"
+            )
+
+            pytree = _replace_node(pytree, j.path, new_value, state_indices)
+
+    return pytree
+
+
+def autoconvert(
+    pytree: PyTree, state_dict: dict, pytree_model_order: list[str] | None = None
+) -> PyTree:
+    torchfields = state_dict_to_fields(state_dict)
+    jaxfields, state_indices = pytree_to_fields(pytree, pytree_model_order)
+
+    pytree = convert(
+        state_dict,
+        pytree,
+        jaxfields,
+        state_indices,
+        torchfields,
+    )
+
+    return pytree

{statedict2pytree-0.6.0.dist-info → statedict2pytree-1.0.1.dist-info}/METADATA

@@ -1,58 +1,41 @@
 Metadata-Version: 2.3
 Name: statedict2pytree
-Version: 0.6.0
+Version: 1.0.1
 Summary: Converts torch models into PyTrees for Equinox
 Author-email: "Artur A. Galstyan" <mail@arturgalstyan.dev>
-Requires-Python: ~=3.10
-Requires-Dist: anthropic
+Requires-Python: >=3.11
 Requires-Dist: beartype
 Requires-Dist: equinox
-Requires-Dist: flask
 Requires-Dist: jax
 Requires-Dist: jaxlib
-Requires-Dist: jaxonmodels
 Requires-Dist: jaxtyping
-Requires-Dist: loguru
-Requires-Dist: penzai
 Requires-Dist: pydantic
-Requires-Dist: pytest
-Requires-Dist: python-dotenv
-Requires-Dist: torch
-Requires-Dist: torchvision
-Requires-Dist: typing-extensions
+Requires-Dist: tqdm
 Provides-Extra: dev
 Requires-Dist: mkdocs; extra == 'dev'
-Requires-Dist: nox; extra == 'dev'
 Requires-Dist: pre-commit; extra == 'dev'
 Requires-Dist: pytest; extra == 'dev'
+Requires-Dist: torch; extra == 'dev'
 Provides-Extra: examples
 Requires-Dist: jaxonmodels; extra == 'examples'
 Description-Content-Type: text/markdown
 
 # statedict2pytree
 
-![statedict2pytree](statedict2pytree.png "A ResNet demo")
 
-## Docs
+## Update:
 
-Docs can be found [here](https://artur-galstyan.github.io/statedict2pytree/).
+For examples for `statedict2pytree`, check out my other repository [jaxonmodels](https://github.com/Artur-Galstyan/jaxonmodels).
 
+## Docs
 
-## Important
+Docs can be found [here](https://artur-galstyan.github.io/statedict2pytree/).
 
-This package is still in its infancy and hihgly experimental! The code works, but it's far from perfect. With more and more iterations, it will eventually become stable and well tested.
-PRs and other contributions are *highly* welcome! :)
 
 ## Info
 
-`statedict2pytree` is a powerful tool for converting PyTorch state dictionaries to JAX pytrees. It provides both programmatic and UI-based methods for mapping between PyTorch and JAX model parameters.
+`statedict2pytree` is a powerful tool for converting PyTorch state dictionaries to JAX pytrees, specifically for Equinox
 
-## Features
-
-- Convert PyTorch statedicts to JAX pytrees
-- Handle large models with chunked file conversion
-- Provide an "intuitive-ish" UI for parameter mapping
-- Support both in-memory and file-based conversions
 
 ## Installation
 
@@ -64,7 +47,6 @@ The goal of this package is to simplify the conversion from PyTorch models into
 
 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?
 
@@ -73,8 +55,8 @@ Currently, there is no sophisticated shape matching in place. Two matrices are c
 (8, 1, 1) and (8, ) match, because (8 _ 1 _ 1 = 8)
 
 
-
 ### Disclaimer
 
 Some of the docstrings and the docs have been written with the help of
 Claude.
+

statedict2pytree-1.0.1.dist-info/RECORD

@@ -0,0 +1,5 @@
+statedict2pytree/__init__.py,sha256=nOtp-ZeTwqHaIibPrqWWXJYzB5PnZHhaRw0HSr9fLBk,215
+statedict2pytree/converter.py,sha256=Toehra2_guidbR5Gyl1uUi1YYv37jkoSKDQjxo1tzwg,8626
+statedict2pytree-1.0.1.dist-info/METADATA,sha256=Z_tRmXwxVMGUMrOYm3SiCxnl4L9OUINZMrBqxbEg5dE,1918
+statedict2pytree-1.0.1.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
+statedict2pytree-1.0.1.dist-info/RECORD,,

client/.gitignore

@@ -1,3 +0,0 @@
-.DS_Store
-node_modules
-public/bundle.*