PyPI - statedict2pytree - Versions diffs - 0.1.2__py3-none-any.whl - Mend

statedict2pytree 0.1.2__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Files changed (8) hide show

statedict2pytree/__init__.py +4 -0
statedict2pytree/statedict2pytree.py +194 -0
statedict2pytree/static/input.css +3 -0
statedict2pytree/static/output.css +1734 -0
statedict2pytree/templates/index.html +246 -0
statedict2pytree-0.1.2.dist-info/METADATA +43 -0
statedict2pytree-0.1.2.dist-info/RECORD +8 -0
statedict2pytree-0.1.2.dist-info/WHEEL +4 -0

statedict2pytree/__init__.py ADDED Viewed

@@ -0,0 +1,4 @@
+from statedict2pytree.statedict2pytree import (
+    convert as convert,
+    start_conversion as start_conversion,
+)

statedict2pytree/statedict2pytree.py ADDED Viewed

@@ -0,0 +1,194 @@
+import functools as ft
+import re
+import equinox as eqx
+import flask
+import jax
+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 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)
+    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)
+    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 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
+    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):
+        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],
+                )
+    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)

statedict2pytree/static/input.css ADDED Viewed

@@ -0,0 +1,3 @@
+@tailwind base;
+@tailwind components;
+@tailwind utilities;