statedict2pytree 0.4.0__py3-none-any.whl → 0.5.2__py3-none-any.whl

client/src/App.svelte

@@ -0,0 +1,343 @@
+<script lang="ts">
+    //@ts-ignore
+    import Sortable from "sortablejs/modular/sortable.complete.esm.js";
+    import { onMount } from "svelte";
+    import Swal from "sweetalert2";
+
+    let model: string = "model.eqx";
+
+    const Toast = Swal.mixin({
+        toast: true,
+        position: "top-end",
+        showConfirmButton: false,
+        timer: 5000,
+        timerProgressBar: true,
+        didOpen: (toast) => {
+            toast.onmouseenter = Swal.stopTimer;
+            toast.onmouseleave = Swal.resumeTimer;
+        },
+    });
+    type Field = {
+        path: string;
+        shape: Number[];
+        skip: boolean;
+        type: string | null;
+    };
+
+    let jaxFields: Field[] = [];
+    let torchFields: Field[] = [];
+    let torchSortable: Sortable;
+    onMount(async () => {
+        let req = await fetch("/startup/getJaxFields");
+        jaxFields = (await req.json()) as Field[];
+        req = await fetch("/startup/getTorchFields");
+        torchFields = await req.json();
+        setTimeout(() => {
+            initSortable();
+        }, 100);
+    });
+
+    function initSortable() {
+        torchSortable = new Sortable(document.getElementById("torch-fields"), {
+            animation: 150,
+            multiDrag: true,
+            ghostClass: "bg-blue-400",
+            selectedClass: "bg-accent",
+            multiDragKey: "shift",
+            onEnd: onEnd,
+        });
+    }
+
+    function swap(a: any, b: any, array: any[]) {
+        const temp = array[a];
+        array[a] = array[b];
+        array[b] = temp;
+        return array;
+    }
+
+    function fetchJaxAndTorchFields() {
+        let allTorchElements =
+            document.getElementById("torch-fields")?.children;
+        if (!allTorchElements) {
+            Toast.fire({
+                icon: "error",
+                title: "Couldn't find PyTorch elements",
+            });
+            return {
+                error: "Failed to fetch PyTorch elements",
+                jaxFields: [],
+                torchFields: [],
+            };
+        }
+        let allTorchFields: HTMLElement[] = [];
+        for (let i = 0; i < allTorchElements.length; i++) {
+            allTorchFields.push(allTorchElements[i].firstChild as HTMLElement);
+        }
+
+        let newTorchFields: Field[] = [];
+        allTorchFields.forEach((el) => {
+            newTorchFields.push({
+                path: el.getAttribute("data-path"),
+                shape: el
+                    .getAttribute("data-shape")
+                    ?.split(",")
+                    .map((x) => parseInt(x)),
+                type: el.getAttribute("data-type"),
+                skip: el.getAttribute("data-skip") === "true",
+            } as Field);
+        });
+
+        const allJaxFields = document.querySelectorAll('[data-jax="jax"]');
+        let newJaxFields: Field[] = [];
+        allJaxFields.forEach((el) => {
+            newJaxFields.push({
+                path: el.getAttribute("data-path"),
+                shape: el
+                    .getAttribute("data-shape")
+                    ?.split(",")
+                    .map((x) => parseInt(x)),
+                type: el.getAttribute("data-type"),
+                skip: el.getAttribute("data-skip") === "true",
+            } as Field);
+        });
+
+        return { jaxFields: newJaxFields, torchFields: newTorchFields };
+    }
+
+    function onEnd() {
+        const updatedFields = fetchJaxAndTorchFields();
+        if (updatedFields.error) {
+            Toast.fire({
+                icon: "error",
+                title: updatedFields.error,
+            });
+            return;
+        }
+
+        for (let i = 0; i < updatedFields.jaxFields.length; i++) {
+            let jaxField = updatedFields.jaxFields[i];
+            let torchField = updatedFields.torchFields[i];
+            if (torchField === undefined) continue;
+            if (torchField.skip === true) {
+                document
+                    .getElementById("jax-" + i)
+                    ?.classList.remove("bg-error");
+                document
+                    .getElementById("torch-" + i)
+                    ?.classList.remove("bg-error");
+                continue;
+            }
+            let jaxShape = jaxField.shape;
+            let torchShape = torchField.shape;
+            //@ts-ignore
+            let jaxShapeProduct = jaxShape.reduce((a, b) => a * b, 1);
+            //@ts-ignore
+            let torchShapeProduct = torchShape.reduce((a, b) => a * b, 1);
+            if (jaxShapeProduct !== torchShapeProduct) {
+                document.getElementById("jax-" + i)?.classList.add("bg-error");
+                document
+                    .getElementById("torch-" + i)
+                    ?.classList.add("bg-error");
+            } else {
+                document
+                    .getElementById("jax-" + i)
+                    ?.classList.remove("bg-error");
+                document
+                    .getElementById("torch-" + i)
+                    ?.classList.remove("bg-error");
+            }
+        }
+
+        if (updatedFields.torchFields.length > updatedFields.jaxFields.length) {
+            for (
+                let i = updatedFields.jaxFields.length;
+                i < updatedFields.torchFields.length;
+                i++
+            ) {
+                document
+                    .getElementById("torch-" + i)
+                    ?.classList.remove("bg-error");
+            }
+        }
+    }
+    function checkFields(jaxFields: Field[], torchFields: Field[]) {
+        if (jaxFields.length > torchFields.length) {
+            return {
+                error: "JAX and PyTorch have lengths! Make sure to pad the PyTorch side.",
+            };
+        }
+
+        for (let i = 0; i < jaxFields.length; i++) {
+            let jaxField = jaxFields[i];
+            let torchField = torchFields[i];
+            if (torchField.skip === true) {
+                continue;
+            }
+
+            //@ts-ignore
+            let jaxShapeProduct = jaxField.shape.reduce((a, b) => a * b, 1);
+            //@ts-ignore
+            let torchShapeProduct = torchField.shape.reduce((a, b) => a * b, 1);
+
+            if (jaxShapeProduct !== torchShapeProduct) {
+                return {
+                    error: `JAX ${jaxField.path} with shape ${jaxField.shape} doesn't match PyTorch ${torchField.path} with shape ${torchField.shape}`,
+                };
+            }
+        }
+        return { success: true };
+    }
+    function removeSkipLayer(index: number) {
+        torchFields = torchFields.toSpliced(index, 1);
+        setTimeout(() => {
+            onEnd();
+        }, 100);
+    }
+    function addSkipLayer(index: number) {
+        const newField = {
+            skip: true,
+            shape: [],
+            path: "",
+            type: "",
+        } as Field;
+        torchFields = torchFields.toSpliced(index, 0, newField);
+        setTimeout(() => {
+            onEnd();
+        }, 100);
+    }
+
+    async function convert() {
+        let fields = fetchJaxAndTorchFields();
+        if (fields.error) {
+            Toast.fire({
+                icon: "error",
+                title: fields.error,
+            });
+            return;
+        }
+
+        let check = checkFields(fields.jaxFields, fields.torchFields);
+        if (check.error) {
+            Toast.fire({
+                icon: "error",
+                title: "Failed to convert",
+                text: check.error,
+            });
+            return;
+        }
+
+        const response = await fetch("/convert", {
+            method: "POST",
+            headers: {
+                "Content-Type": "application/json",
+            },
+            body: JSON.stringify({
+                model: model,
+                jaxFields: fields.jaxFields,
+                torchFields: fields.torchFields,
+            }),
+        });
+
+        const res = await response.json();
+        console.log(res);
+        if (res.error) {
+            Toast.fire({
+                icon: "error",
+                title: res.error,
+            });
+        } else {
+            Toast.fire({
+                icon: "success",
+                title: "Conversion successful",
+            });
+        }
+    }
+</script>
+
+<svelte:head><title>Statedict2PyTree</title></svelte:head>
+
+<h1 class="text-3xl my-12">Welcome to Torch2Jax</h1>
+
+<div class="grid grid-cols-2 gap-x-2">
+    <div class="">
+        <h2 class="text-2xl">JAX</h2>
+        <div id="jax-fields" class="">
+            {#each jaxFields as field, i}
+                <div
+                    class="border h-12 rounded-xl flex flex-col justify-center"
+                >
+                    <div
+                        id={"jax-" + String(i)}
+                        class="whitespace-nowrap overflow-x-scroll cursor-pointer mx-2"
+                        data-jax="jax"
+                        data-path={field.path}
+                        data-shape={field.shape}
+                        data-skip={field.skip}
+                        data-type={field.type}
+                    >
+                        {field.path}
+                        {field.shape}
+                    </div>
+                </div>
+            {/each}
+        </div>
+    </div>
+
+    <div class="">
+        <h2 class="text-2xl">PyTorch</h2>
+        <div id="torch-fields" class="">
+            {#each torchFields as field, i}
+                <div class="flex space-x-2 border h-12 rounded-xl">
+                    <div
+                        id={"torch-" + String(i)}
+                        data-torch="torch"
+                        data-path={field.path}
+                        data-shape={field.shape}
+                        data-skip={field.skip}
+                        data-type={field.type}
+                        class="flex-1 mx-2 my-auto whitespace-nowrap overflow-x-scroll cursor-pointer"
+                    >
+                        {#if field.skip}
+                            SKIP
+                        {:else}
+                            {field.path}
+                            {field.shape}
+                        {/if}
+                    </div>
+                    {#if field.skip}
+                        <button
+                            class="btn btn-ghost"
+                            on:click={() => {
+                                removeSkipLayer(i);
+                            }}>-</button
+                        >
+                    {/if}
+                    <button
+                        class="btn btn-ghost"
+                        on:click={() => {
+                            addSkipLayer(i);
+                        }}>+</button
+                    >
+                </div>
+            {/each}
+        </div>
+    </div>
+</div>
+<div class="flex justify-center my-12 w-full">
+    <div class="flex flex-col justify-center w-full">
+        <input
+            id="name"
+            type="text"
+            name="name"
+            class="input input-primary w-full"
+            placeholder="Name of the new file (model.eqx per default)"
+            bind:value={model}
+        />
+        <button
+            on:click={convert}
+            class="btn btn-accent btn-wide btn-lg mx-auto my-2"
+        >
+            Convert!
+        </button>
+    </div>
+</div>

client/src/main.js

@@ -0,0 +1,8 @@
+import App from "./App.svelte";
+
+const app = new App({
+  target: document.body,
+  props: {},
+});
+
+export default app;

client/tailwind.config.js

@@ -0,0 +1,8 @@
+/** @type {import('tailwindcss').Config} */
+module.exports = {
+  content: ["./src/**/*.{html,js,svelte}"],
+  theme: {
+    extend: {},
+  },
+  plugins: [require("daisyui")],
+};

client/tsconfig.json

@@ -0,0 +1,5 @@
+{
+  "extends": "@tsconfig/svelte/tsconfig.json",
+  "include": ["src/**/*", "src/node_modules"],
+  "exclude": ["node_modules/*", "__sapper__/*", "public/*"]
+}

statedict2pytree-0.5.2.dist-info/METADATA

@@ -0,0 +1,64 @@
+Metadata-Version: 2.3
+Name: statedict2pytree
+Version: 0.5.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
+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: torch
+Requires-Dist: torchvision
+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'
+Provides-Extra: examples
+Requires-Dist: jaxonmodels; extra == 'examples'
+Description-Content-Type: text/markdown
+
+# statedict2pytree
+
+![statedict2pytree](statedict2pytree.png "A ResNet demo")
+
+## Important
+
+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
+
+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:
+
+(8, 1, 1) and (8, ) match, because (8 _ 1 _ 1 = 8)
+
+## Get Started
+
+### Installation
+
+Run
+
+```bash
+pip install statedict2pytree
+```
+
+### Docs
+
+Documentation will appear as soon as I have all the necessary features implemented. Until then, check out the "main.py" file for a better example.

statedict2pytree-0.5.2.dist-info/RECORD

@@ -0,0 +1,17 @@
+client/.gitignore,sha256=dUAd7J1wJVSaU9WD0Gp4c7_5yEz5lOdWLfsjubHf1z8,39
+client/package-lock.json,sha256=JCW5mq9bGNs29yGO1EIlcysDaEuTm2i_AJoyt-45yO8,158531
+client/package.json,sha256=Ad-MDEQeh7BPHWPYLd3u9sXk8YVuO_dXmpkxxU1Pglo,1044
+client/rollup.config.mjs,sha256=RAepJhL2V5Rf-BlJBZJxllVl0mxtr67GSVr9aU0JUnA,1073
+client/tailwind.config.js,sha256=TfN5eOoFOUPGBou6OoK54M14PtokgxWDJUsV4qkurS8,175
+client/tsconfig.json,sha256=cLHEFXx-Q55XqbF9QjQ4XScSEQ15n-vS5tsTcqY4UAY,158
+client/public/bundle.js,sha256=l4Vu_8_7v6k7XFisnI8jCFzyMeztr9jc3Z3lrrPDpk0,347197
+client/public/bundle.js.map,sha256=s9zOkP-34BWrQFunVbkouwTKVAfL8EyIUcfKgSicH8M,682783
+client/public/index.html,sha256=jUx-NPKkFN2EF2lj-8Ml49CEHxKJFWK9seszauI4GE0,335
+client/public/input.css,sha256=zBp60NAZ3bHTLQ7LWIugrCbOQdhiXdbDZjSLJfg6KOw,59
+client/public/output.css,sha256=3iiBiTGfqAeVKuRZRqgcixX3ztSnlp0zqHXkjSKtmVs,38664
+client/src/App.svelte,sha256=hHVoQ_C2xGMmd4d86ftZFeTlGmOcOW_wJ6abq0_qvWo,11170
+client/src/empty.ts,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+client/src/main.js,sha256=O_8UgVd1vJM8BcHO7U_6jkL76ZSA6oC7GLLcL9F3JLA,118
+statedict2pytree-0.5.2.dist-info/METADATA,sha256=YgE7bgWDMI6urA71bH-zWleR_mw6SJ5QPneUZVvHL2E,2242
+statedict2pytree-0.5.2.dist-info/WHEEL,sha256=zEMcRr9Kr03x1ozGwg5v9NQBKn3kndp6LSoSlVg-jhU,87
+statedict2pytree-0.5.2.dist-info/RECORD,,

statedict2pytree/__init__.py

@@ -1,7 +0,0 @@
-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

@@ -1,218 +0,0 @@
-import functools as ft
-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
-from penzai import pz
-from pydantic import BaseModel
-
-
-app = flask.Flask(__name__)
-
-
-class Field(BaseModel):
-    path: str
-    shape: tuple[int, ...]
-
-
-class TorchField(Field):
-    pass
-
-
-class JaxField(Field):
-    type: str
-
-
-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:
-    if len(targets) == 0 or tree is None:
-        return tree
-    else:
-        next_target: str = targets[0]
-        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]
-            if hasattr(tree, name):
-                subtree = getattr(tree, name)[int(index)]
-            else:
-                subtree = None
-                if log_when_not_found:
-                    logger.info(f"Couldn't find  {name} in {tree.__class__}")
-        else:
-            if hasattr(tree, next_target):
-                subtree = getattr(tree, next_target)
-            else:
-                subtree = None
-                if log_when_not_found:
-                    logger.info(f"Couldn't find  {next_target} in {tree.__class__}")
-        return get_node(subtree, targets[1:])
-
-
-def pytree_to_fields(pytree: PyTree) -> list[JaxField]:
-    flattened, _ = jax.tree_util.tree_flatten_with_path(pytree)
-    fields: list[JaxField] = []
-    for key_path, value in flattened:
-        path = jax.tree_util.keystr(key_path)
-        type_path = path.split(".")[1:-1]
-        target_path = path.split(".")[1:]
-        node_type = type(get_node(pytree, type_path, log_when_not_found=True))
-        node = get_node(pytree, target_path, log_when_not_found=True)
-        if node is not None and hasattr(node, "shape") and len(node.shape) > 0:
-            fields.append(
-                JaxField(path=path, type=str(node_type), shape=tuple(node.shape))
-            )
-
-    return fields
-
-
-def state_dict_to_fields(state_dict: Optional[dict]) -> 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
-
-
-@app.route("/visualize", methods=["POST"])
-def visualize_with_penzai():
-    global PYTREE, STATE_DICT
-    if PYTREE is None or STATE_DICT is None:
-        return flask.jsonify({"error": "No Pytree or StateDict found"})
-    request_data = flask.request.json
-    if request_data is None:
-        return flask.jsonify({"error": "No data received"})
-    jax_fields = request_data["jaxFields"]
-    torch_fields = request_data["torchFields"]
-    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)
-
-    combined_html = f"<html><body>{html_jax}<hr>{html_torch}</body></html>"
-    return combined_html
-
-
-@app.route("/convert", methods=["POST"])
-def convert_torch_to_jax():
-    global PYTREE, STATE_DICT
-    if PYTREE is None or STATE_DICT is None:
-        return flask.jsonify({"error": "No Pytree or StateDict found"})
-    request_data = flask.request.json
-    if request_data is None:
-        return flask.jsonify({"error": "No data received"})
-
-    jax_fields_json = request_data["jaxFields"]
-    jax_fields: list[JaxField] = []
-    for f in jax_fields_json:
-        shape_tuple = tuple(
-            [int(i) for i in f["shape"].strip("()").split(",") if len(i) > 0]
-        )
-        jax_fields.append(JaxField(path=f["path"], type=f["type"], shape=shape_tuple))
-
-    torch_fields_json = request_data["torchFields"]
-    torch_fields: list[TorchField] = []
-    for f in torch_fields_json:
-        shape_tuple = tuple(
-            [int(i) for i in f["shape"].strip("()").split(",") if len(i) > 0]
-        )
-        torch_fields.append(TorchField(path=f["path"], shape=shape_tuple))
-
-    name = request_data["name"]
-    model, state = convert(jax_fields, torch_fields, PYTREE, STATE_DICT)
-    eqx.tree_serialise_leaves(name, (model, state))
-
-    return flask.jsonify({"status": "success"})
-
-
-@app.route("/", methods=["GET"])
-def index():
-    pytree_fields = pytree_to_fields(PYTREE)
-    return flask.render_template(
-        "index.html",
-        pytree_fields=pytree_fields,
-        torch_fields=state_dict_to_fields(STATE_DICT),
-    )
-
-
-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 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))
-
-        else:
-            where = ft.partial(get_node, targets=path)
-            if where(model) is not None:
-                model = eqx.tree_at(
-                    where,
-                    model,
-                    state_dict[torch_field.path].reshape(jax_field.shape),
-                )
-    result: dict[str, list[TorchField]] = {}
-    for tuple_item in state_paths:
-        path_prefix = tuple_item[0].path.split(".")[1:-1]
-        prefix_key = ".".join(path_prefix)
-        if prefix_key not in result:
-            result[prefix_key] = []
-        result[prefix_key].append(tuple_item[1])
-
-    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]])
-            state = state.set(state_index, to_replace_tuple)
-    return model, state
-
-
-def start_conversion(pytree: PyTree, state_dict: dict):
-    global PYTREE, STATE_DICT
-    if state_dict is None:
-        raise ValueError("STATE_DICT must not be None!")
-    PYTREE = pytree
-    STATE_DICT = state_dict
-
-    for k, v in STATE_DICT.items():
-        STATE_DICT[k] = v.numpy()
-    app.run(debug=True, port=5500)