statedict2pytree 0.6.0__tar.gz → 1.0.0__tar.gz

statedict2pytree-1.0.0/.github/workflows/run_tests.yml

@@ -0,0 +1,38 @@
+name: Run tests
+
+on: [pull_request]
+
+jobs:
+  run-test:
+    strategy:
+      matrix:
+        python-version: ["3.10"]
+        os: [ubuntu-latest]
+      fail-fast: false
+    runs-on: ${{ matrix.os }}
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v1
+
+      - name: Create virtual environment and install dependencies
+        run: |
+          uv venv --python ${{ matrix.python-version }}
+          source .venv/bin/activate
+          uv pip install -e ".[dev]"
+          uv pip install -e .
+
+      - name: Run pre-commit hooks
+        uses: pre-commit/action@v3.0.1
+
+      - name: Run tests
+        run: |
+          source .venv/bin/activate
+          python3 -m pytest

statedict2pytree-1.0.0/.pre-commit-config.yaml

@@ -0,0 +1,12 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.11.12
+    hooks:
+      - id: ruff-check
+      - id: ruff-format
+  - repo: https://github.com/RobertCraigie/pyright-python
+    rev: v1.1.351
+    hooks:
+      - id: pyright
+        additional_dependencies:
+          [beartype, jax, jaxtyping, pytest, typing_extensions]

{statedict2pytree-0.6.0 → statedict2pytree-1.0.0}/PKG-INFO

@@ -1,58 +1,41 @@
 Metadata-Version: 2.3
 Name: statedict2pytree
-Version: 0.6.0
+Version: 1.0.0
 Summary: Converts torch models into PyTrees for Equinox
 Author-email: "Artur A. Galstyan" <mail@arturgalstyan.dev>
 Requires-Python: ~=3.10
-Requires-Dist: anthropic
 Requires-Dist: beartype
 Requires-Dist: equinox
-Requires-Dist: flask
-Requires-Dist: jax
+Requires-Dist: jax>=0.6.1
 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-0.6.0 → statedict2pytree-1.0.0}/README.md

@@ -1,27 +1,19 @@
 # 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
 
@@ -33,7 +25,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?
 
@@ -42,8 +33,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.0/docs/index.md

@@ -0,0 +1,186 @@
+# Quickstart Guide
+
+## Installation
+
+To install `statedict2pytree`, run:
+
+```bash
+pip install statedict2pytree
+```
+
+---
+
+## Basic Usage
+
+There are 4-5 main functions you might interact with:
+
+* `autoconvert`
+* `convert`
+* `pytree_to_fields`
+* `state_dict_to_fields`
+* `move_running_fields_to_the_end` (optional helper)
+
+---
+
+## General Information
+
+`statedict2pytree` primarily aligns your JAX PyTree and the PyTorch `state_dict` side-by-side. It then checks if the shapes of the aligned weights match. If they do, it converts the PyTorch tensors to JAX arrays and places them into a new PyTree with the same structure as your original JAX PyTree.
+
+**This means that the order and the shape of the arrays in your PyTree and the `state_dict` must match after any optional reordering!** The `pytree_to_fields` function uses a filter (defaulting to `equinox.is_array`) to determine which elements are considered fields.
+
+For example, this conversion will **work** ✅:
+
+| Parameter         | JAX Shape    | PyTorch Shape |
+| :---------------- | :----------- | :------------ |
+| `linear.weight`   | `(2, 2)`     | `(2, 2)`      |
+| `linear.bias`     | `(2,)`       | `(2,)`        |
+| `conv.weight`     | `(1, 1, 2, 2)` | `(1, 1, 2, 2)`|
+| `conv.bias`       | `(1,)`       | `(1,)`        |
+
+Since the shapes match when aligned in the same order, the conversion is successful.
+
+On the other hand, this will **not work** ❌:
+
+| Parameter         | JAX Shape    | PyTorch Shape | Mismatch? |
+| :---------------- | :----------- | :------------ | :-------- |
+| `linear.weight`   | `(2, 2)`     | `(3, 2)`      | Yes       |
+| `linear.bias`     | `(2,)`       | `(3,)`        | Yes       |
+| `conv.weight`     | `(1, 1, 2, 2)` | `(1, 1, 2, 2)`| No        |
+| `conv.bias`       | `(1,)`       | `(1,)`        | No        |
+
+This conversion will fail because the shapes of `model.linear.weight` and `model.linear.bias` don't match between the PyTree and the state dict.
+
+Another reason why the conversion might fail is if the **order** of parameters (and thus the shapes of misaligned parameters) doesn't match:
+
+| JAX Parameter (Model Order) | JAX Shape      | PyTorch Counterpart (`state_dict` Order) | PyTorch Shape  | Issue if Matched Sequentially                     |
+| :-------------------------- | :------------- | :--------------------------------------- | :------------- | :------------------------------------------------ |
+| `model['conv']['weight']`   | `(1, 1, 2, 2)` | `state_dict['model.linear.weight']`      | `(2, 2)`       | Order: JAX `conv.w` `(1122)` vs PT `linear.w` `(22)` |
+| `model['conv']['bias']`     | `(1,)`         | `state_dict['model.linear.bias']`        | `(2,)`         | Order: JAX `conv.b` `(1,)` vs PT `linear.b` `(2,)`   |
+| `model['linear']['weight']` | `(2, 2)`       | `state_dict['model.conv.weight']`        | `(1, 1, 2, 2)` | Order: JAX `linear.w` `(22)` vs PT `conv.w` `(1122)`|
+| `model['linear']['bias']`   | `(2,)`         | `state_dict['model.conv.bias']`          | `(1,)`         | Order: JAX `linear.b` `(2,)` vs PT `conv.b` `(1,)`   |
+
+To help with the order issue, you can provide a `list[str]` specifying the desired order of PyTree fields (matching the `state_dict`'s conceptual order, or vice-versa if you reorder `state_dict` fields). This is especially helpful when you can't easily force the correct order using `move_running_fields_to_the_end`. For the example above, if your PyTree expects `conv` then `linear`, the list of strings representing the *names from the state\_dict in the JAX PyTree's desired order* would be:
+
+```python
+['model.conv.weight', 'model.conv.bias', 'model.linear.weight', 'model.linear.bias']
+```
+This list would be passed to `pytree_to_fields` via `autoconvert`'s `pytree_model_order` argument to ensure `jaxfields` are in this sequence. Alternatively, you could reorder `torchfields` using `move_running_fields_to_the_end` or other custom logic.
+
+---
+
+## API Reference
+
+### `autoconvert`
+
+This is the simplest, highest-level function for most use cases.
+
+```python
+def autoconvert(
+    pytree: PyTree,
+    state_dict: dict,
+    pytree_model_order: list[str] | None = None
+) -> PyTree:
+    ...
+```
+
+You provide your JAX `pytree` and the PyTorch `state_dict`. Optionally, you can give `pytree_model_order` (a list of strings representing `jax.tree_util.keystr(path)`) to ensure the JAX fields are processed in a specific sequence. It handles the steps of field extraction (using `pytree_to_fields` with its default `filter=eqx.is_array`), alignment, and conversion, returning the populated JAX PyTree. If you need custom filtering for PyTree leaves, you should use `pytree_to_fields` and `convert` separately.
+
+* **Parameters**:
+    * `pytree`: The JAX PyTree (e.g., an Equinox model) whose structure is the target.
+    * `state_dict`: The PyTorch state dictionary containing the weights.
+    * `pytree_model_order` (optional): A list of JAX KeyPath strings (like `'.layers.0.linear.weight'`). If provided, JAX fields will be ordered according to this list. This is useful if the automatic PyTree traversal order doesn't match the `state_dict` order.
+* **Returns**: A new JAX PyTree with the same structure as the input `pytree`, but with weights populated from the `state_dict`.
+
+---
+
+### `convert`
+
+This is the core function that performs the actual conversion once the JAX PyTree fields and PyTorch `state_dict` fields have been extracted and aligned.
+
+```python
+def convert(
+    state_dict: dict[str, Any],
+    pytree: PyTree,
+    jaxfields: list[JaxField],
+    state_indices: dict | None,
+    torchfields: list[TorchField],
+    dtype: Any | None = None,
+) -> PyTree:
+    ...
+```
+
+It iterates through the aligned `jaxfields` and `torchfields`, checks for shape compatibility (reshapability), converts PyTorch tensors (expected as values in `state_dict`) to JAX arrays (optionally casting `dtype`), and inserts them into the correct place in the JAX PyTree.
+
+* **Parameters**:
+    * `state_dict`: The original PyTorch state dictionary. Values are expected to be tensor-like (e.g., `torch.Tensor`).
+    * `pytree`: The JAX PyTree that will be populated.
+    * `jaxfields`: An ordered list of `JaxField` objects (obtained from `pytree_to_fields`) representing the leaves of the JAX PyTree.
+    * `state_indices`: A dictionary mapping state markers to `eqx.nn.StateIndex` objects, used for handling Equinox stateful layers.
+    * `torchfields`: An ordered list of `TorchField` objects (obtained from `state_dict_to_fields`) representing the tensors in the PyTorch `state_dict`. **This list must be ordered to match `jaxfields`**.
+    * `dtype` (optional): The JAX data type to convert floating-point tensors to (e.g., `jnp.float32`). Defaults to JAX's current default floating-point type.
+* **Returns**: A new JAX PyTree populated with weights from the `state_dict`.
+
+---
+
+### `pytree_to_fields`
+
+This function traverses a JAX PyTree and extracts information about its array leaves based on a filter.
+
+```python
+def pytree_to_fields(
+    pytree: PyTree,
+    model_order: list[str] | None = None,
+    filter: Callable[[Array], bool] = eqx.is_array,
+) -> tuple[list[JaxField], dict | None]:
+    ...
+```
+
+It identifies all JAX arrays (or other elements satisfying the `filter`) within the `pytree`, recording their `KeyPath` (path within the PyTree) and shape. If `model_order` is provided, it attempts to reorder the extracted fields according to that list. This is crucial for ensuring the JAX fields align correctly with the PyTorch fields.
+
+* **Parameters**:
+    * `pytree`: The JAX PyTree to analyze.
+    * `model_order` (optional): A list of strings, where each string is a `jax.tree_util.keystr` representation of a `KeyPath` to an array leaf in the `pytree`. If provided, the output `JaxField` list will be sorted according to this order, with any fields not in `model_order` appended at the end.
+    * `filter` (optional): A callable that takes a PyTree leaf (e.g., an array) and returns `True` if it should be considered a field to be converted, `False` otherwise. Defaults to `equinox.is_array`.
+* **Returns**: A tuple containing:
+    * `list[JaxField]`: A list of `JaxField` objects, each describing a filtered leaf in the PyTree (path, shape).
+    * `dict | None`: A dictionary containing information about `eqx.nn.StateIndex` objects found in the PyTree, or `None` if none are found.
+
+---
+
+### `state_dict_to_fields`
+
+This function processes a PyTorch `state_dict` to extract information about its tensors.
+
+```python
+def state_dict_to_fields(
+    state_dict: dict[str, Any],
+) -> list[TorchField]:
+    ...
+```
+
+It iterates through the `state_dict`, creating a `TorchField` object for each value that has a `shape` attribute and a non-empty shape (typically tensors). This object stores the tensor's name (key in the `state_dict`) and its shape.
+
+* **Parameters**:
+    * `state_dict`: The PyTorch state dictionary. Values are typically `torch.Tensor` or other array-like objects.
+* **Returns**: A list of `TorchField` objects, each describing a tensor in the `state_dict` (path/key, shape). The order matches the iteration order of the input `state_dict`.
+
+---
+
+### `move_running_fields_to_the_end`
+
+This is an optional utility function to help reorder fields extracted from a PyTorch `state_dict`.
+
+```python
+def move_running_fields_to_the_end(
+    torchfields: list[TorchField],
+    identifier: str = "running_"
+):
+    ...
+```
+
+It's particularly useful for models with layers like `BatchNorm`, where PyTorch often stores `running_mean` and `running_var` interspersed with weights and biases, while Equinox (a common JAX library) typically expects stateful components like these at the end of a layer's parameter list. This function moves any `TorchField` whose path contains the `identifier` (defaulting to `"running_"`) to the end of the list.
+
+* **Parameters**:
+    * `torchfields`: The list of `TorchField` objects to be reordered.
+    * `identifier` (optional): A string that, if found within a `TorchField`'s path, will cause that field to be moved to the end of the list. Default is `"running_"`.
+* **Returns**: The modified list of `TorchField` objects with identified fields moved to the end.

statedict2pytree-1.0.0/mkdocs.yml

@@ -0,0 +1,19 @@
+site_name: Statedict2Pytree
+
+theme:
+  name: material
+
+plugins:
+  - search
+  - mkdocstrings:
+      default_handler: python
+
+nav:
+  - Home: index.md
+
+markdown_extensions:
+  - pymdownx.highlight:
+      anchor_linenums: true
+  - pymdownx.inlinehilite
+  - pymdownx.snippets
+  - pymdownx.superfences

statedict2pytree-1.0.0/pyproject.toml

@@ -0,0 +1,24 @@
+[project]
+name = "statedict2pytree"
+version = "1.0.0"
+description = "Converts torch models into PyTrees for Equinox"
+readme = "README.md"
+requires-python = "~=3.10"
+authors = [{ name = "Artur A. Galstyan", email = "mail@arturgalstyan.dev" }]
+dependencies = [
+    "jax>=0.6.1",
+    "equinox",
+    "jaxlib",
+    "beartype",
+    "jaxtyping",
+    "pydantic",
+    "tqdm",
+]
+
+[project.optional-dependencies]
+dev = ["pre-commit", "pytest", "mkdocs", "torch"]
+examples = ["jaxonmodels"]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"

statedict2pytree-1.0.0/pyrightconfig.json

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

statedict2pytree-1.0.0/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