statedict2pytree 0.3.0__tar.gz → 0.5.0__tar.gz

{statedict2pytree-0.3.0 → statedict2pytree-0.5.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: statedict2pytree
-Version: 0.3.0
+Version: 0.5.0
 Summary: Converts torch models into PyTrees for Equinox
 Author-email: "Artur A. Galstyan" <mail@arturgalstyan.dev>
 Requires-Python: ~=3.10

statedict2pytree-0.5.0/examples/convert_resnet.py

@@ -0,0 +1,19 @@
+import jax
+import statedict2pytree as s2p
+from resnet import resnet18
+from torchvision.models import resnet18 as t_resnet18, ResNet18_Weights
+
+
+def convert_resnet():
+    resnet_jax = resnet18(key=jax.random.PRNGKey(33), make_with_state=False)
+    resnet_torch = t_resnet18(weights=ResNet18_Weights.DEFAULT)
+    state_dict = resnet_torch.state_dict()
+
+    s2p.start_conversion(resnet_jax, state_dict)
+    # model, state = s2p.autoconvert(resnet_jax, state_dict)
+    # name = "resnet18.eqx"
+    # eqx.tree_serialise_leaves(name, (model, state))
+
+
+if __name__ == "__main__":
+    convert_resnet()

{statedict2pytree-0.3.0 → statedict2pytree-0.5.0}/examples/resnet.py

@@ -333,12 +333,22 @@ def resnet18(
 
 
 def resnet34(
-    image_channels: int = 3, num_classes: int = 1000, *, key: PRNGKeyArray, **kwargs
+    image_channels: int = 3,
+    num_classes: int = 1000,
+    *,
+    key: PRNGKeyArray,
+    make_with_state: bool = True,
+    **kwargs,
 ):
     layers = [3, 4, 6, 3]
-    return eqx.nn.make_with_state(ResNet)(
-        BasicBlock, layers, image_channels, num_classes, **kwargs, key=key
-    )
+    if make_with_state:
+        return eqx.nn.make_with_state(ResNet)(
+            BasicBlock, layers, image_channels, num_classes, **kwargs, key=key
+        )
+    else:
+        return ResNet(
+            BasicBlock, layers, image_channels, num_classes, **kwargs, key=key
+        )
 
 
 def resnet50(
@@ -361,18 +371,39 @@ def resnet50(
 
 
 def resnet101(
-    image_channels: int = 3, num_classes: int = 1000, *, key: PRNGKeyArray, **kwargs
+    image_channels: int = 3,
+    num_classes: int = 1000,
+    *,
+    key: PRNGKeyArray,
+    make_with_state: bool = True,
+    **kwargs,
 ):
     layers = [3, 4, 23, 3]
-    return eqx.nn.make_with_state(ResNet)(
-        Bottleneck, layers, image_channels, num_classes, **kwargs, key=key
-    )
+
+    if make_with_state:
+        return eqx.nn.make_with_state(ResNet)(
+            Bottleneck, layers, image_channels, num_classes, **kwargs, key=key
+        )
+    else:
+        return ResNet(
+            Bottleneck, layers, image_channels, num_classes, **kwargs, key=key
+        )
 
 
 def resnet152(
-    image_channels: int = 3, num_classes: int = 1000, *, key: PRNGKeyArray, **kwargs
+    image_channels: int = 3,
+    num_classes: int = 1000,
+    *,
+    key: PRNGKeyArray,
+    make_with_state: bool = True,
+    **kwargs,
 ):
     layers = [3, 8, 36, 3]
-    return eqx.nn.make_with_state(ResNet)(
-        Bottleneck, layers, image_channels, num_classes, **kwargs, key=key
-    )
+    if make_with_state:
+        return eqx.nn.make_with_state(ResNet)(
+            Bottleneck, layers, image_channels, num_classes, **kwargs, key=key
+        )
+    else:
+        return ResNet(
+            Bottleneck, layers, image_channels, num_classes, **kwargs, key=key
+        )

statedict2pytree-0.3.0/examples/test_resnet_inference.py → statedict2pytree-0.5.0/examples/resnet_inference.py

@@ -6,15 +6,15 @@ import equinox as eqx
 import jax
 import jax.numpy as jnp
 import torch
+from examples.resnet import resnet18
 from PIL import Image
-from tests.resnet import resnet50
 from torchvision import transforms
-from torchvision.models import resnet50 as t_resnet50, ResNet50_Weights
+from torchvision.models import resnet18 as t_resnet18, ResNet18_Weights
 
 
 def test_resnet():
-    resnet_jax = resnet50(key=jax.random.PRNGKey(33), make_with_state=False)
-    resnet_torch = t_resnet50(weights=ResNet50_Weights.DEFAULT)
+    resnet_jax = resnet18(key=jax.random.PRNGKey(33), make_with_state=False)
+    resnet_torch = t_resnet18(weights=ResNet18_Weights.DEFAULT)
 
     img_name = "doggo.jpeg"
 
@@ -42,7 +42,7 @@ def test_resnet():
     )  # Outputs the ImageNet class index of the prediction
 
     url = "https://storage.googleapis.com/download.tensorflow.org/data/imagenet_class_index.json"
-    with urllib.request.urlopen(url) as url:
+    with urllib.request.urlopen(url) as url:  # pyright: ignore
         imagenet_labels = json.loads(url.read().decode())
 
     label = imagenet_labels[str(predicted.item())][1]

statedict2pytree-0.5.0/package.json

@@ -0,0 +1,9 @@
+{
+  "scripts": {
+    "watch": "npx tailwindcss -i ./statedict2pytree/static/input.css -o ./statedict2pytree/static/output.css --watch"
+  },
+  "devDependencies": {
+    "daisyui": "^4.11.1",
+    "tailwindcss": "^3.4.3"
+  }
+}

{statedict2pytree-0.3.0 → statedict2pytree-0.5.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "statedict2pytree"
-version = "0.3.0"
+version = "0.5.0"
 description = "Converts torch models into PyTrees for Equinox"
 readme = "README.md"
 requires-python = "~=3.10"

{statedict2pytree-0.3.0 → statedict2pytree-0.5.0}/statedict2pytree/statedict2pytree.py

@@ -46,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]
@@ -157,6 +157,9 @@ 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)
 
 
@@ -212,4 +215,5 @@ def start_conversion(pytree: PyTree, state_dict: dict):
 
     for k, v in STATE_DICT.items():
         STATE_DICT[k] = v.numpy()
+    app.jinja_env.globals.update(enumerate=enumerate)
     app.run(debug=True, port=5500)

{statedict2pytree-0.3.0 → statedict2pytree-0.5.0}/statedict2pytree/static/output.css

@@ -784,23 +784,42 @@ html {
   }
 }
 
-.avatar.placeholder > div {
-  display: flex;
+.alert {
+  display: grid;
+  width: 100%;
+  grid-auto-flow: row;
+  align-content: flex-start;
   align-items: center;
-  justify-content: center;
+  justify-items: center;
+  gap: 1rem;
+  text-align: center;
+  border-radius: var(--rounded-box, 1rem);
+  border-width: 1px;
+  --tw-border-opacity: 1;
+  border-color: var(--fallback-b2,oklch(var(--b2)/var(--tw-border-opacity)));
+  padding: 1rem;
+  --tw-text-opacity: 1;
+  color: var(--fallback-bc,oklch(var(--bc)/var(--tw-text-opacity)));
+  --alert-bg: var(--fallback-b2,oklch(var(--b2)/1));
+  --alert-bg-mix: var(--fallback-b1,oklch(var(--b1)/1));
+  background-color: var(--alert-bg);
 }
 
-@media (hover:hover) {
-  .menu li > *:not(ul, .menu-title, details, .btn):active,
-.menu li > *:not(ul, .menu-title, details, .btn).active,
-.menu li > details > summary:active {
-    --tw-bg-opacity: 1;
-    background-color: var(--fallback-n,oklch(var(--n)/var(--tw-bg-opacity)));
-    --tw-text-opacity: 1;
-    color: var(--fallback-nc,oklch(var(--nc)/var(--tw-text-opacity)));
+@media (min-width: 640px) {
+  .alert {
+    grid-auto-flow: column;
+    grid-template-columns: auto minmax(auto,1fr);
+    justify-items: start;
+    text-align: start;
   }
 }
 
+.avatar.placeholder > div {
+  display: flex;
+  align-items: center;
+  justify-content: center;
+}
+
 .btn {
   display: inline-flex;
   height: 3rem;
@@ -934,18 +953,6 @@ html {
       border-color: color-mix(in oklab, var(--fallback-p,oklch(var(--p)/1)) 90%, black);
     }
   }
-
-  :where(.menu li:not(.menu-title, .disabled) > *:not(ul, details, .menu-title)):not(.active, .btn):hover, :where(.menu li:not(.menu-title, .disabled) > details > summary:not(.menu-title)):not(.active, .btn):hover {
-    cursor: pointer;
-    outline: 2px solid transparent;
-    outline-offset: 2px;
-  }
-
-  @supports (color: oklch(0% 0 0)) {
-    :where(.menu li:not(.menu-title, .disabled) > *:not(ul, details, .menu-title)):not(.active, .btn):hover, :where(.menu li:not(.menu-title, .disabled) > details > summary:not(.menu-title)):not(.active, .btn):hover {
-      background-color: var(--fallback-bc,oklch(var(--bc)/0.1));
-    }
-  }
 }
 
 .input {
@@ -978,59 +985,6 @@ html {
   text-decoration-line: underline;
 }
 
-.menu {
-  display: flex;
-  flex-direction: column;
-  flex-wrap: wrap;
-  font-size: 0.875rem;
-  line-height: 1.25rem;
-  padding: 0.5rem;
-}
-
-.menu :where(li ul) {
-  position: relative;
-  white-space: nowrap;
-  margin-inline-start: 1rem;
-  padding-inline-start: 0.5rem;
-}
-
-.menu :where(li:not(.menu-title) > *:not(ul, details, .menu-title, .btn)), .menu :where(li:not(.menu-title) > details > summary:not(.menu-title)) {
-  display: grid;
-  grid-auto-flow: column;
-  align-content: flex-start;
-  align-items: center;
-  gap: 0.5rem;
-  grid-auto-columns: minmax(auto, max-content) auto max-content;
-  -webkit-user-select: none;
-     -moz-user-select: none;
-          user-select: none;
-}
-
-.menu li.disabled {
-  cursor: not-allowed;
-  -webkit-user-select: none;
-     -moz-user-select: none;
-          user-select: none;
-  color: var(--fallback-bc,oklch(var(--bc)/0.3));
-}
-
-.menu :where(li > .menu-dropdown:not(.menu-dropdown-show)) {
-  display: none;
-}
-
-:where(.menu li) {
-  position: relative;
-  display: flex;
-  flex-shrink: 0;
-  flex-direction: column;
-  flex-wrap: wrap;
-  align-items: stretch;
-}
-
-:where(.menu li) .badge {
-  justify-self: end;
-}
-
 .toast {
   position: fixed;
   display: flex;
@@ -1042,6 +996,14 @@ html {
   padding: 1rem;
 }
 
+.alert-error {
+  border-color: var(--fallback-er,oklch(var(--er)/0.2));
+  --tw-text-opacity: 1;
+  color: var(--fallback-erc,oklch(var(--erc)/var(--tw-text-opacity)));
+  --alert-bg: var(--fallback-er,oklch(var(--er)/1));
+  --alert-bg-mix: var(--fallback-b1,oklch(var(--b1)/1));
+}
+
 @media (prefers-reduced-motion: no-preference) {
   .btn {
     animation: button-pop var(--animation-btn, 0.25s) ease-out;
@@ -1273,88 +1235,6 @@ html {
   outline-offset: 2px;
 }
 
-:where(.menu li:empty) {
-  --tw-bg-opacity: 1;
-  background-color: var(--fallback-bc,oklch(var(--bc)/var(--tw-bg-opacity)));
-  opacity: 0.1;
-  margin: 0.5rem 1rem;
-  height: 1px;
-}
-
-.menu :where(li ul):before {
-  position: absolute;
-  bottom: 0.75rem;
-  inset-inline-start: 0px;
-  top: 0.75rem;
-  width: 1px;
-  --tw-bg-opacity: 1;
-  background-color: var(--fallback-bc,oklch(var(--bc)/var(--tw-bg-opacity)));
-  opacity: 0.1;
-  content: "";
-}
-
-.menu :where(li:not(.menu-title) > *:not(ul, details, .menu-title, .btn)),
-.menu :where(li:not(.menu-title) > details > summary:not(.menu-title)) {
-  border-radius: var(--rounded-btn, 0.5rem);
-  padding-left: 1rem;
-  padding-right: 1rem;
-  padding-top: 0.5rem;
-  padding-bottom: 0.5rem;
-  text-align: start;
-  transition-property: color, background-color, border-color, text-decoration-color, fill, stroke, opacity, box-shadow, transform, filter, -webkit-backdrop-filter;
-  transition-property: color, background-color, border-color, text-decoration-color, fill, stroke, opacity, box-shadow, transform, filter, backdrop-filter;
-  transition-property: color, background-color, border-color, text-decoration-color, fill, stroke, opacity, box-shadow, transform, filter, backdrop-filter, -webkit-backdrop-filter;
-  transition-timing-function: cubic-bezier(0.4, 0, 0.2, 1);
-  transition-timing-function: cubic-bezier(0, 0, 0.2, 1);
-  transition-duration: 200ms;
-  text-wrap: balance;
-}
-
-:where(.menu li:not(.menu-title, .disabled) > *:not(ul, details, .menu-title)):not(summary, .active, .btn).focus, :where(.menu li:not(.menu-title, .disabled) > *:not(ul, details, .menu-title)):not(summary, .active, .btn):focus, :where(.menu li:not(.menu-title, .disabled) > *:not(ul, details, .menu-title)):is(summary):not(.active, .btn):focus-visible, :where(.menu li:not(.menu-title, .disabled) > details > summary:not(.menu-title)):not(summary, .active, .btn).focus, :where(.menu li:not(.menu-title, .disabled) > details > summary:not(.menu-title)):not(summary, .active, .btn):focus, :where(.menu li:not(.menu-title, .disabled) > details > summary:not(.menu-title)):is(summary):not(.active, .btn):focus-visible {
-  cursor: pointer;
-  background-color: var(--fallback-bc,oklch(var(--bc)/0.1));
-  --tw-text-opacity: 1;
-  color: var(--fallback-bc,oklch(var(--bc)/var(--tw-text-opacity)));
-  outline: 2px solid transparent;
-  outline-offset: 2px;
-}
-
-.menu li > *:not(ul, .menu-title, details, .btn):active,
-.menu li > *:not(ul, .menu-title, details, .btn).active,
-.menu li > details > summary:active {
-  --tw-bg-opacity: 1;
-  background-color: var(--fallback-n,oklch(var(--n)/var(--tw-bg-opacity)));
-  --tw-text-opacity: 1;
-  color: var(--fallback-nc,oklch(var(--nc)/var(--tw-text-opacity)));
-}
-
-.menu :where(li > details > summary)::-webkit-details-marker {
-  display: none;
-}
-
-.menu :where(li > details > summary):after,
-.menu :where(li > .menu-dropdown-toggle):after {
-  justify-self: end;
-  display: block;
-  margin-top: -0.5rem;
-  height: 0.5rem;
-  width: 0.5rem;
-  transform: rotate(45deg);
-  transition-property: transform, margin-top;
-  transition-duration: 0.3s;
-  transition-timing-function: cubic-bezier(0.4, 0, 0.2, 1);
-  content: "";
-  transform-origin: 75% 75%;
-  box-shadow: 2px 2px;
-  pointer-events: none;
-}
-
-.menu :where(li > details[open] > summary):after,
-.menu :where(li > .menu-dropdown-toggle.menu-dropdown-show):after {
-  transform: rotate(225deg);
-  margin-top: 0;
-}
-
 .mockup-browser .mockup-browser-toolbar .input {
   position: relative;
   margin-left: auto;
@@ -1552,115 +1432,6 @@ html {
   transform: translate(var(--tw-translate-x), var(--tw-translate-y)) rotate(var(--tw-rotate)) skewX(var(--tw-skew-x)) skewY(var(--tw-skew-y)) scaleX(var(--tw-scale-x)) scaleY(var(--tw-scale-y));
 }
 
-.tooltip {
-  position: relative;
-  display: inline-block;
-  --tooltip-offset: calc(100% + 1px + var(--tooltip-tail, 0px));
-}
-
-.tooltip:before {
-  position: absolute;
-  pointer-events: none;
-  z-index: 1;
-  content: var(--tw-content);
-  --tw-content: attr(data-tip);
-}
-
-.tooltip:before, .tooltip-top:before {
-  transform: translateX(-50%);
-  top: auto;
-  left: 50%;
-  right: auto;
-  bottom: var(--tooltip-offset);
-}
-
-.tooltip {
-  position: relative;
-  display: inline-block;
-  text-align: center;
-  --tooltip-tail: 0.1875rem;
-  --tooltip-color: var(--fallback-n,oklch(var(--n)/1));
-  --tooltip-text-color: var(--fallback-nc,oklch(var(--nc)/1));
-  --tooltip-tail-offset: calc(100% + 0.0625rem - var(--tooltip-tail));
-}
-
-.tooltip:before,
-.tooltip:after {
-  opacity: 0;
-  transition-property: color, background-color, border-color, text-decoration-color, fill, stroke, opacity, box-shadow, transform, filter, -webkit-backdrop-filter;
-  transition-property: color, background-color, border-color, text-decoration-color, fill, stroke, opacity, box-shadow, transform, filter, backdrop-filter;
-  transition-property: color, background-color, border-color, text-decoration-color, fill, stroke, opacity, box-shadow, transform, filter, backdrop-filter, -webkit-backdrop-filter;
-  transition-delay: 100ms;
-  transition-duration: 200ms;
-  transition-timing-function: cubic-bezier(0.4, 0, 0.2, 1);
-}
-
-.tooltip:after {
-  position: absolute;
-  content: "";
-  border-style: solid;
-  border-width: var(--tooltip-tail, 0);
-  width: 0;
-  height: 0;
-  display: block;
-}
-
-.tooltip:before {
-  max-width: 20rem;
-  border-radius: 0.25rem;
-  padding-left: 0.5rem;
-  padding-right: 0.5rem;
-  padding-top: 0.25rem;
-  padding-bottom: 0.25rem;
-  font-size: 0.875rem;
-  line-height: 1.25rem;
-  background-color: var(--tooltip-color);
-  color: var(--tooltip-text-color);
-  width: -moz-max-content;
-  width: max-content;
-}
-
-.tooltip.tooltip-open:before {
-  opacity: 1;
-  transition-delay: 75ms;
-}
-
-.tooltip.tooltip-open:after {
-  opacity: 1;
-  transition-delay: 75ms;
-}
-
-.tooltip:hover:before {
-  opacity: 1;
-  transition-delay: 75ms;
-}
-
-.tooltip:hover:after {
-  opacity: 1;
-  transition-delay: 75ms;
-}
-
-.tooltip:has(:focus-visible):after,
-.tooltip:has(:focus-visible):before {
-  opacity: 1;
-  transition-delay: 75ms;
-}
-
-.tooltip:not([data-tip]):hover:before,
-.tooltip:not([data-tip]):hover:after {
-  visibility: hidden;
-  opacity: 0;
-}
-
-.tooltip:after, .tooltip-top:after {
-  transform: translateX(-50%);
-  border-color: var(--tooltip-color) transparent transparent transparent;
-  top: auto;
-  left: 50%;
-  right: auto;
-  bottom: var(--tooltip-tail-offset);
-}
-
 .static {
   position: static;
 }
@@ -1684,14 +1455,42 @@ html {
   display: flex;
 }
 
+.grid {
+  display: grid;
+}
+
+.hidden {
+  display: none;
+}
+
+.h-6 {
+  height: 1.5rem;
+}
+
 .w-10\/12 {
   width: 83.333333%;
 }
 
+.w-6 {
+  width: 1.5rem;
+}
+
 .w-full {
   width: 100%;
 }
 
+.shrink-0 {
+  flex-shrink: 0;
+}
+
+.cursor-pointer {
+  cursor: pointer;
+}
+
+.grid-cols-2 {
+  grid-template-columns: repeat(2, minmax(0, 1fr));
+}
+
 .flex-col {
   flex-direction: column;
 }
@@ -1700,14 +1499,17 @@ html {
   justify-content: center;
 }
 
-.space-x-4 > :not([hidden]) ~ :not([hidden]) {
-  --tw-space-x-reverse: 0;
-  margin-right: calc(1rem * var(--tw-space-x-reverse));
-  margin-left: calc(1rem * calc(1 - var(--tw-space-x-reverse)));
+.gap-x-2 {
+  -moz-column-gap: 0.5rem;
+       column-gap: 0.5rem;
 }
 
-.rounded-box {
-  border-radius: var(--rounded-box, 1rem);
+.overflow-x-scroll {
+  overflow-x: scroll;
+}
+
+.whitespace-nowrap {
+  white-space: nowrap;
 }
 
 .bg-base-200 {
@@ -1715,8 +1517,18 @@ html {
   background-color: var(--fallback-b2,oklch(var(--b2)/var(--tw-bg-opacity)));
 }
 
-.text-left {
-  text-align: left;
+.bg-blue-400 {
+  --tw-bg-opacity: 1;
+  background-color: rgb(96 165 250 / var(--tw-bg-opacity));
+}
+
+.bg-error {
+  --tw-bg-opacity: 1;
+  background-color: var(--fallback-er,oklch(var(--er)/var(--tw-bg-opacity)));
+}
+
+.stroke-current {
+  stroke: currentColor;
 }
 
 .text-2xl {
@@ -1728,7 +1540,3 @@ html {
   font-size: 1.875rem;
   line-height: 2.25rem;
 }
-
-.ease-in-out {
-  transition-timing-function: cubic-bezier(0.4, 0, 0.2, 1);
-}

{statedict2pytree-0.3.0 → statedict2pytree-0.5.0}/statedict2pytree/templates/index.html

@@ -11,26 +11,73 @@
       href="{{ url_for('static', filename='output.css') }}"
     />
     <script src="https://cdn.jsdelivr.net/npm/sweetalert2@11"></script>
+    <script src="https://cdn.jsdelivr.net/npm/sortablejs@latest/Sortable.min.js"></script>
     <script type="module">
-      import {
-        Draggable,
-        Sortable,
-        Droppable,
-        Swappable,
-        Plugins,
-      } from "https://cdn.jsdelivr.net/npm/@shopify/draggable/build/esm/index.mjs";
-      const sortableLists = document.querySelectorAll(".draggable-list");
-
-      for (const sortableList of sortableLists) {
-        const sortable = new Sortable(sortableList, {
-          draggable: "li",
-          plugins: [Plugins.SortAnimation],
-          swapAnimation: {
-            duration: 200,
-            easingFunction: "ease-in-out",
+      let jaxSortable = new Sortable(document.getElementById("jax-fields"), {
+        animation: 150,
+        ghostClass: "blue-background-class",
+      });
+
+      let torchSortable = new Sortable(
+        document.getElementById("torch-fields"),
+        {
+          animation: 150,
+          ghostClass: "bg-blue-400",
+          onEnd: function (evt) {
+            document.getElementById("error-field").classList.add("hidden");
+            let allJaxFields = document.querySelectorAll("#jax-fields > div");
+            let allTorchFields = document.querySelectorAll(
+              "#torch-fields > div",
+            );
+            if (allJaxFields.length !== allTorchFields.length) {
+              Swal.fire({
+                icon: "error",
+                title:
+                  "The number of fields in JAX and PyTorch should be the same",
+              });
+            } else {
+              for (let i = 0; i < allJaxFields.length; i++) {
+                let jaxField = allJaxFields[i];
+                let torchField = allTorchFields[i];
+                let jaxShape = jaxField.getAttribute("data-shape");
+                let torchShape = torchField.getAttribute("data-shape");
+
+                jaxShape = jaxShape
+                  .replace("(", "")
+                  .replace(")", "")
+                  .replace(/\s+/g, "")
+                  .replace(/,\s*$/, "");
+
+                torchShape = torchShape
+                  .replace("(", "")
+                  .replace(")", "")
+                  .replace(/\s+/g, "")
+                  .replace(/,\s*$/, "");
+
+                let jaxShapeParts = jaxShape.split(",").map((x) => parseInt(x));
+                let torchShapeParts = torchShape
+                  .split(",")
+                  .map((x) => parseInt(x));
+                let jaxShapeProduct = jaxShapeParts.reduce((a, b) => a * b, 1);
+                let torchShapeProduct = torchShapeParts.reduce(
+                  (a, b) => a * b,
+                  1,
+                );
+
+                let jaxEl = jaxField;
+                let torchEl = torchField;
+                if (jaxShapeProduct !== torchShapeProduct) {
+                  jaxEl.classList.add("bg-error");
+                  torchEl.classList.add("bg-error");
+                } else {
+                  jaxEl.classList.remove("bg-error");
+                  torchEl.classList.remove("bg-error");
+                }
+              }
+            }
           },
-        });
-      }
+        },
+      );
     </script>
   </head>
   <body class="w-10/12 mx-auto">
@@ -66,10 +113,8 @@
         xhr.setRequestHeader("Content-Type", "application/json");
         xhr.onload = function () {
           if (xhr.status >= 200 && xhr.status < 300) {
-            // Successfully received response
             var container = document.getElementById("visualizationResult");
             container.innerHTML = xhr.responseText;
-            // Execute any scripts that were in the response
             var scripts = container.getElementsByTagName("script");
             for (var i = 0; i < scripts.length; i++) {
               var script = document.createElement("script");
@@ -86,7 +131,7 @@
 
       function getJaxAndTorchFields() {
         const jaxFields = Array.from(
-          document.querySelectorAll(".draggable-list")[0].children,
+          document.querySelectorAll("#jax-fields")[0].children,
         ).map((li) => {
           const path = li.getAttribute("data-path");
           const shape = li.getAttribute("data-shape");
@@ -95,7 +140,7 @@
         });
 
         const torchFields = Array.from(
-          document.querySelectorAll(".draggable-list")[1].children,
+          document.querySelectorAll("#torch-fields")[0].children,
         ).map((li) => {
           const path = li.getAttribute("data-path");
           const shape = li.getAttribute("data-shape");
@@ -121,6 +166,11 @@
               icon: "error",
               title: `${jaxFields[i].path} has shape ${jaxFields[i].shape}, while ${torchFields[i].path} has shape ${torchFields[i].shape}`,
             });
+            document.getElementById("error-field").classList.remove("hidden");
+            document
+              .getElementById("error-field")
+              .querySelector("span").innerText =
+              `${jaxFields[i].path} has shape ${jaxFields[i].shape}, while ${torchFields[i].path} has shape ${torchFields[i].shape}`;
             return { error: "Invalid shapes" };
           }
         }
@@ -133,7 +183,8 @@
         if (fields.error) {
           Toast.fire({
             icon: "error",
-            title: "The number of fields in JAX and PyTorch should be the same",
+            title: "Failed to convert!",
+            text: fields.error,
           });
         }
         const jaxFields = fields.jaxFields;
@@ -178,42 +229,36 @@
     </script>
     <h1 class="text-3xl my-12">Welcome to Torch2Jax</h1>
 
-    <div class="flex space-x-4">
-      <div class="w-full">
+    <div class="grid grid-cols-2 gap-x-2">
+      <div class="">
         <h2 class="text-2xl">JAX</h2>
-        <ul class="draggable-list menu bg-base-200 rounded-box">
+        <div id="jax-fields" class="bg-base-200">
           {% for field in pytree_fields %}
-          <li
-            draggable="true"
+          <div
             data-path="{{field.path}}"
             data-shape="{{field.shape}}"
             data-type="{{field.type}}"
+            class="whitespace-nowrap overflow-x-scroll cursor-pointer"
           >
-            <p
-              class="tooltip text-left"
-              data-path="{{field.path}}"
-              data-tip="{{field.type}}"
-            >
-              {{ field.path}} {{field.shape }}
-            </p>
-          </li>
+            {{ field.path }} {{ field.shape }}
+          </div>
           {% endfor %}
-        </ul>
+        </div>
       </div>
 
-      <div class="w-full">
+      <div class="">
         <h2 class="text-2xl">PyTorch</h2>
-        <ul class="draggable-list menu bg-base-200 rounded-box w-full">
+        <div id="torch-fields" class="bg-base-200">
           {% for field in torch_fields %}
-          <li
-            draggable="true"
+          <div
             data-path="{{field.path}}"
             data-shape="{{field.shape}}"
+            class="whitespace-nowrap overflow-x-scroll cursor-pointer"
           >
-            <p class="text-left">{{ field.path }} {{ field.shape }}</p>
-          </li>
+            {{ field.path }} {{ field.shape }}
+          </div>
           {% endfor %}
-        </ul>
+        </div>
       </div>
     </div>
     <div class="flex justify-center my-12 w-full">
@@ -234,6 +279,23 @@
         </button>
       </div>
     </div>
+    <div role="alert" class="alert alert-error hidden" id="error-field">
+      <svg
+        xmlns="http://www.w3.org/2000/svg"
+        class="stroke-current shrink-0 h-6 w-6"
+        fill="none"
+        viewBox="0 0 24 24"
+      >
+        <path
+          stroke-linecap="round"
+          stroke-linejoin="round"
+          stroke-width="2"
+          d="M10 14l2-2m0 0l2-2m-2 2l-2-2m2 2l2 2m7-2a9 9 0 11-18 0 9 9 0 0118 0z"
+        />
+      </svg>
+      <span></span>
+    </div>
+
     <div class="flex justify-center">
       <button onclick="visualize()" class="btn btn-secondary">
         Visualize with Penzai!

{statedict2pytree-0.3.0 → statedict2pytree-0.5.0}/tailwind.config.js

@@ -1,6 +1,6 @@
 /** @type {import('tailwindcss').Config} */
 module.exports = {
-  content: ["./torch2jax/**/*.{html,js}"],
+  content: ["./statedict2pytree/**/*.{html,js}"],
   theme: {
     extend: {},
   },

statedict2pytree-0.5.0/tests/test_batchnorm.py

@@ -0,0 +1,46 @@
+import equinox as eqx
+import jax
+import numpy as np
+import statedict2pytree as s2p
+import torch
+
+
+def test_linear():
+    in_features = 10
+    out_features = 10
+
+    class J(eqx.Module):
+        linear: eqx.nn.Linear
+        norm: eqx.nn.BatchNorm
+
+        def __init__(self):
+            self.linear = eqx.nn.Linear(
+                in_features, out_features, key=jax.random.PRNGKey(30)
+            )
+            self.norm = eqx.nn.BatchNorm(input_size=out_features, axis_name="batch")
+
+    class T(torch.nn.Module):
+        def __init__(self) -> None:
+            super(T, self).__init__()
+            self.linear = torch.nn.Linear(in_features, out_features)
+            self.norm = torch.nn.BatchNorm1d(out_features)
+
+    jax_model = J()
+    torch_model = T()
+    state_dict = torch_model.state_dict()
+
+    model, state = s2p.autoconvert(pytree=jax_model, state_dict=state_dict)
+
+    assert np.allclose(
+        np.array(model.linear.weight), torch_model.linear.weight.detach().numpy()
+    )
+    assert np.allclose(
+        np.array(model.linear.bias), torch_model.linear.bias.detach().numpy()
+    )
+
+    assert np.allclose(
+        np.array(model.norm.weight), torch_model.norm.weight.detach().numpy()
+    )
+    assert np.allclose(
+        np.array(model.norm.bias), torch_model.norm.bias.detach().numpy()
+    )

{statedict2pytree-0.3.0 → statedict2pytree-0.5.0}/tests/test_conv.py

@@ -55,7 +55,3 @@ def test_conv():
             np.array(model.conv.bias),
             torch_model.conv.bias.detach().numpy().reshape(model.conv.bias.shape),
         )
-
-
-if __name__ == "__main__":
-    test_conv()

{statedict2pytree-0.3.0 → statedict2pytree-0.5.0}/tests/test_linear.py

@@ -39,7 +39,3 @@ def test_linear():
     assert np.allclose(
         np.array(model.linear.bias), torch_model.linear.bias.detach().numpy()
     )
-
-
-if __name__ == "__main__":
-    test_linear()

statedict2pytree-0.3.0/examples/convert_resnet.py

@@ -1,16 +0,0 @@
-import jax
-import statedict2pytree as s2p
-from resnet import resnet50
-from torchvision.models import resnet50 as t_resnet50, ResNet50_Weights
-
-
-def convert_resnet():
-    resnet_jax = resnet50(key=jax.random.PRNGKey(33), make_with_state=False)
-    resnet_torch = t_resnet50(weights=ResNet50_Weights.DEFAULT)
-    state_dict = resnet_torch.state_dict()
-
-    s2p.start_conversion(resnet_jax, state_dict)
-
-
-if __name__ == "__main__":
-    convert_resnet()

statedict2pytree-0.3.0/package.json

@@ -1,10 +0,0 @@
-{
-  "scripts": {
-    "dev": "python3 torch2jax/torch2jax.py --debug=True",
-    "watch": "npx tailwindcss -i ./torch2jax/static/input.css -o ./torch2jax/static/output.css --watch"
-  },
-  "devDependencies": {
-    "daisyui": "^4.11.1",
-    "tailwindcss": "^3.4.3"
-  }
-}