Merge pull request #114 from biggraph/darabos-repeats

.github/workflows/test.yaml

@@ -81,6 +81,7 @@ jobs:
       - name: Run Playwright tests
         run: |
           cd lynxkite-app/web
+          npm run build
           npm run test
 
       - uses: actions/upload-artifact@v4

examples/Model definition

@@ -1,171 +1,278 @@
 {
   "edges": [
     {
-      "id": "Input: embedding 1 Linear 1",
-      "source": "Input: embedding 1",
-      "sourceHandle": "x",
-      "target": "Linear 1",
+      "id": "MSE loss 2 Optimizer 2",
+      "source": "MSE loss 2",
+      "sourceHandle": "output",
+      "target": "Optimizer 2",
+      "targetHandle": "loss"
+    },
+    {
+      "id": "Activation 1 MSE loss 2",
+      "source": "Activation 1",
+      "sourceHandle": "output",
+      "target": "MSE loss 2",
       "targetHandle": "x"
     },
     {
-      "id": "Input: label 1 MSE loss 1",
-      "source": "Input: label 1",
-      "sourceHandle": "y",
-      "target": "MSE loss 1",
+      "id": "Input: tensor 3 MSE loss 2",
+      "source": "Input: tensor 3",
+      "sourceHandle": "x",
+      "target": "MSE loss 2",
       "targetHandle": "y"
     },
     {
-      "id": "Linear 1 Activation 2",
-      "source": "Linear 1",
+      "id": "Activation 1 Repeat 1",
+      "source": "Activation 1",
+      "sourceHandle": "output",
+      "target": "Repeat 1",
+      "targetHandle": "input"
+    },
+    {
+      "id": "Input: tensor 1 Linear 1",
+      "source": "Input: tensor 1",
       "sourceHandle": "x",
-      "target": "Activation 2",
+      "target": "Linear 1",
       "targetHandle": "x"
     },
     {
-      "id": "Activation 2 MSE loss 1",
-      "source": "Activation 2",
-      "sourceHandle": "x",
-      "target": "MSE loss 1",
+      "id": "Linear 1 Activation 1",
+      "source": "Linear 1",
+      "sourceHandle": "output",
+      "target": "Activation 1",
       "targetHandle": "x"
     },
     {
-      "id": "MSE loss 1 Optimizer 2",
-      "source": "MSE loss 1",
-      "sourceHandle": "loss",
-      "target": "Optimizer 2",
-      "targetHandle": "loss"
+      "id": "Repeat 1 Linear 1",
+      "source": "Repeat 1",
+      "sourceHandle": "output",
+      "target": "Linear 1",
+      "targetHandle": "x"
     }
   ],
   "env": "PyTorch model",
   "nodes": [
     {
       "data": {
+        "__execution_delay": 0.0,
+        "collapsed": null,
         "display": null,
         "error": null,
+        "input_metadata": null,
         "meta": {
-          "inputs": {},
-          "name": "Input: embedding",
-          "outputs": {
-            "x": {
-              "name": "x",
-              "position": "top",
+          "inputs": {
+            "loss": {
+              "name": "loss",
+              "position": "bottom",
               "type": {
                 "type": "tensor"
               }
             }
           },
-          "params": {},
+          "name": "Optimizer",
+          "outputs": {},
+          "params": {
+            "lr": {
+              "default": 0.001,
+              "name": "lr",
+              "type": {
+                "type": "<class 'float'>"
+              }
+            },
+            "type": {
+              "default": "AdamW",
+              "name": "type",
+              "type": {
+                "enum": [
+                  "AdamW",
+                  "Adafactor",
+                  "Adagrad",
+                  "SGD",
+                  "Lion",
+                  "Paged AdamW",
+                  "Galore AdamW"
+                ]
+              }
+            }
+          },
           "type": "basic"
         },
-        "params": {},
+        "params": {
+          "lr": "0.1",
+          "type": "SGD"
+        },
         "status": "planned",
-        "title": "Input: embedding"
+        "title": "Optimizer"
       },
       "dragHandle": ".bg-primary",
-      "height": 200.0,
-      "id": "Input: embedding 1",
+      "height": 250.0,
+      "id": "Optimizer 2",
       "position": {
-        "x": 91.0,
-        "y": 266.0
+        "x": 292.3983313429414,
+        "y": -853.8015246037802
       },
       "type": "basic",
-      "width": 200.0
+      "width": 232.0
     },
     {
       "data": {
+        "__execution_delay": 0.0,
+        "collapsed": null,
         "display": null,
         "error": null,
+        "input_metadata": null,
         "meta": {
           "inputs": {
             "x": {
               "name": "x",
               "position": "bottom",
               "type": {
-                "type": "tensor"
+                "type": "<class 'inspect._empty'>"
               }
             }
           },
-          "name": "Linear",
+          "name": "Activation",
           "outputs": {
-            "x": {
-              "name": "x",
+            "output": {
+              "name": "output",
               "position": "top",
               "type": {
-                "type": "tensor"
+                "type": "None"
               }
             }
           },
           "params": {
-            "output_dim": {
-              "default": "same",
-              "name": "output_dim",
+            "type": {
+              "default": "ReLU",
+              "name": "type",
               "type": {
-                "type": "<class 'str'>"
+                "enum": [
+                  "ReLU",
+                  "Leaky_ReLU",
+                  "Tanh",
+                  "Mish"
+                ]
               }
             }
           },
+          "position": {
+            "x": 344.0,
+            "y": 384.0
+          },
           "type": "basic"
         },
         "params": {
-          "output_dim": "same"
+          "type": "Leaky_ReLU"
         },
         "status": "planned",
-        "title": "Linear"
+        "title": "Activation"
       },
       "dragHandle": ".bg-primary",
       "height": 200.0,
-      "id": "Linear 1",
+      "id": "Activation 1",
       "position": {
-        "x": 86.0,
-        "y": 33.0
+        "x": 99.77615018185415,
+        "y": -249.43925929074078
       },
       "type": "basic",
       "width": 200.0
     },
     {
       "data": {
+        "__execution_delay": 0.0,
+        "collapsed": null,
         "display": null,
         "error": null,
+        "input_metadata": null,
         "meta": {
-          "inputs": {
+          "inputs": {},
+          "name": "Input: tensor",
+          "outputs": {
             "x": {
               "name": "x",
-              "position": "bottom",
+              "position": "top",
               "type": {
                 "type": "tensor"
               }
-            },
-            "y": {
-              "name": "y",
-              "position": "bottom",
+            }
+          },
+          "params": {
+            "name": {
+              "default": null,
+              "name": "name",
               "type": {
-                "type": "tensor"
+                "type": "None"
               }
             }
           },
-          "name": "MSE loss",
+          "position": {
+            "x": 258.0,
+            "y": 397.0
+          },
+          "type": "basic"
+        },
+        "params": {
+          "name": "X"
+        },
+        "status": "planned",
+        "title": "Input: tensor"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 200.0,
+      "id": "Input: tensor 1",
+      "position": {
+        "x": 85.83561484252238,
+        "y": 293.6278596776366
+      },
+      "type": "basic",
+      "width": 200.0
+    },
+    {
+      "data": {
+        "__execution_delay": 0.0,
+        "collapsed": null,
+        "display": null,
+        "error": null,
+        "input_metadata": null,
+        "meta": {
+          "inputs": {},
+          "name": "Input: tensor",
           "outputs": {
-            "loss": {
-              "name": "loss",
+            "x": {
+              "name": "x",
               "position": "top",
               "type": {
                 "type": "tensor"
               }
             }
           },
-          "params": {},
+          "params": {
+            "name": {
+              "default": null,
+              "name": "name",
+              "type": {
+                "type": "None"
+              }
+            }
+          },
+          "position": {
+            "x": 1169.0,
+            "y": 340.0
+          },
           "type": "basic"
         },
-        "params": {},
+        "params": {
+          "name": "Y"
+        },
         "status": "planned",
-        "title": "MSE loss"
+        "title": "Input: tensor"
       },
       "dragHandle": ".bg-primary",
       "height": 200.0,
-      "id": "MSE loss 1",
+      "id": "Input: tensor 3",
       "position": {
-        "x": 315.0,
-        "y": -510.0
+        "x": 485.8840220312055,
+        "y": -149.86223034126274
       },
       "type": "basic",
       "width": 200.0
@@ -174,31 +281,51 @@
       "data": {
         "display": null,
         "error": null,
+        "input_metadata": null,
         "meta": {
-          "inputs": {},
-          "name": "Input: label",
-          "outputs": {
+          "inputs": {
+            "x": {
+              "name": "x",
+              "position": "bottom",
+              "type": {
+                "type": "<class 'inspect._empty'>"
+              }
+            },
             "y": {
               "name": "y",
+              "position": "bottom",
+              "type": {
+                "type": "<class 'inspect._empty'>"
+              }
+            }
+          },
+          "name": "MSE loss",
+          "outputs": {
+            "output": {
+              "name": "output",
               "position": "top",
               "type": {
-                "type": "tensor"
+                "type": "None"
               }
             }
           },
           "params": {},
+          "position": {
+            "x": 937.0,
+            "y": 270.0
+          },
           "type": "basic"
         },
         "params": {},
         "status": "planned",
-        "title": "Input: label"
+        "title": "MSE loss"
       },
       "dragHandle": ".bg-primary",
       "height": 200.0,
-      "id": "Input: label 1",
+      "id": "MSE loss 2",
       "position": {
-        "x": 615.0,
-        "y": -165.0
+        "x": 309.4422414664647,
+        "y": -552.1056805642488
       },
       "type": "basic",
       "width": 200.0
@@ -209,58 +336,62 @@
         "collapsed": null,
         "display": null,
         "error": null,
+        "input_metadata": null,
         "meta": {
           "inputs": {
-            "x": {
-              "name": "x",
-              "position": "bottom",
+            "input": {
+              "name": "input",
+              "position": "top",
               "type": {
                 "type": "tensor"
               }
             }
           },
-          "name": "Activation",
+          "name": "Repeat",
           "outputs": {
-            "x": {
-              "name": "x",
-              "position": "top",
+            "output": {
+              "name": "output",
+              "position": "bottom",
               "type": {
                 "type": "tensor"
               }
             }
           },
           "params": {
-            "type": {
-              "default": "ReLU",
-              "name": "type",
+            "same_weights": {
+              "default": false,
+              "name": "same_weights",
               "type": {
-                "enum": [
-                  "ReLU",
-                  "Leaky ReLU",
-                  "Tanh",
-                  "Mish"
-                ]
+                "type": "<class 'bool'>"
+              }
+            },
+            "times": {
+              "default": 1.0,
+              "name": "times",
+              "type": {
+                "type": "<class 'int'>"
               }
             }
           },
           "position": {
-            "x": 419.0,
-            "y": 396.0
+            "x": 487.0,
+            "y": 443.0
           },
           "type": "basic"
         },
         "params": {
-          "type": "Leaky ReLU"
+          "same_weights": false,
+          "times": "2"
         },
         "status": "planned",
-        "title": "Activation"
+        "title": "Repeat"
       },
       "dragHandle": ".bg-primary",
       "height": 200.0,
-      "id": "Activation 2",
+      "id": "Repeat 1",
       "position": {
-        "x": 93.61643829835265,
-        "y": -229.04087132886406
+        "x": -210.0,
+        "y": -135.0
       },
       "type": "basic",
       "width": 200.0
@@ -271,61 +402,54 @@
         "collapsed": null,
         "display": null,
         "error": null,
+        "input_metadata": null,
         "meta": {
           "inputs": {
-            "loss": {
-              "name": "loss",
+            "x": {
+              "name": "x",
               "position": "bottom",
               "type": {
-                "type": "tensor"
+                "type": "<class 'inspect._empty'>"
               }
             }
           },
-          "name": "Optimizer",
-          "outputs": {},
-          "params": {
-            "lr": {
-              "default": 0.001,
-              "name": "lr",
+          "name": "Linear",
+          "outputs": {
+            "output": {
+              "name": "output",
+              "position": "top",
               "type": {
-                "type": "<class 'float'>"
+                "type": "None"
               }
-            },
-            "type": {
-              "default": "AdamW",
-              "name": "type",
+            }
+          },
+          "params": {
+            "output_dim": {
+              "default": 1024.0,
+              "name": "output_dim",
               "type": {
-                "enum": [
-                  "AdamW",
-                  "Adafactor",
-                  "Adagrad",
-                  "SGD",
-                  "Lion",
-                  "Paged AdamW",
-                  "Galore AdamW"
-                ]
+                "type": "<class 'int'>"
               }
             }
           },
           "position": {
-            "x": 526.0,
-            "y": 116.0
+            "x": 359.0,
+            "y": 310.0
           },
           "type": "basic"
         },
         "params": {
-          "lr": "0.1",
-          "type": "SGD"
+          "output_dim": "4"
         },
         "status": "planned",
-        "title": "Optimizer"
+        "title": "Linear"
       },
       "dragHandle": ".bg-primary",
       "height": 200.0,
-      "id": "Optimizer 2",
+      "id": "Linear 1",
       "position": {
-        "x": 305.6132943499785,
-        "y": -804.0094318451224
+        "x": 88.83370222907377,
+        "y": 48.642890099180136
       },
       "type": "basic",
       "width": 200.0

examples/Model use

@@ -575,58 +575,58 @@
               "columns": [
                 "x",
                 "y",
-                "predicted"
+                "pred"
               ],
               "data": [
                 [
-                  "[0.49691743 0.61873293 0.90698647 0.94486356]",
-                  "[1.49691749 1.61873293 1.90698647 1.94486356]",
-                  "[1.4993021488189697, 1.6404846906661987, 1.923316240310669, 1.9422152042388916]"
+                  "[0.19908059 0.17570406 0.51475513 0.1893943 ]",
+                  "[1.19908059 1.175704   1.51475513 1.18939424]",
+                  "[1.560641884803772, 1.5941988229751587, 1.5775359869003296, 1.4935821294784546]"
                 ],
                 [
-                  "[0.56922203 0.98222166 0.76851749 0.28615737]",
-                  "[1.56922197 1.9822216  1.76851749 1.28615737]",
-                  "[1.5835213661193848, 1.9884355068206787, 1.7694181203842163, 1.2917503118515015]"
+                  "[0.43681622 0.74680805 0.83598751 0.12414402]",
+                  "[1.43681622 1.74680805 1.83598757 1.12414408]",
+                  "[1.5766589641571045, 1.7117265462875366, 1.7645087242126465, 1.3384637832641602]"
                 ],
                 [
-                  "[0.90817457 0.89270043 0.38583666 0.66566533]",
-                  "[1.90817451 1.89270043 1.3858366  1.66566539]",
-                  "[1.9053494930267334, 1.9083378314971924, 1.3998609781265259, 1.6636812686920166]"
+                  "[0.9829582  0.59269661 0.40120947 0.95487177]",
+                  "[1.9829582  1.59269667 1.40120947 1.95487177]",
+                  "[1.5375217199325562, 1.4159281253814697, 1.2972962856292725, 1.7269455194473267]"
                 ],
                 [
-                  "[0.72795159 0.79317838 0.27832931 0.96576637]",
-                  "[1.72795153 1.79317832 1.27832937 1.96576643]",
-                  "[1.734963297843933, 1.8026459217071533, 1.2926064729690552, 1.9596911668777466]"
+                  "[0.32565445 0.90939188 0.07488042 0.13730896]",
+                  "[1.32565451 1.90939188 1.07488036 1.13730896]",
+                  "[1.562728762626648, 1.6061222553253174, 1.597141146659851, 1.4772177934646606]"
                 ],
                 [
-                  "[0.04508126 0.76880038 0.80721325 0.62542385]",
-                  "[1.04508126 1.76880038 1.80721331 1.62542391]",
-                  "[1.0830243825912476, 1.7584562301635742, 1.8005754947662354, 1.6277496814727783]"
+                  "[0.31518555 0.49643308 0.11509258 0.95458382]",
+                  "[1.31518555 1.49643302 1.11509252 1.95458388]",
+                  "[1.528311848640442, 1.3380011320114136, 1.171952247619629, 1.8305948972702026]"
                 ],
                 [
-                  "[0.6032477  0.83361369 0.18538666 0.19108021]",
-                  "[1.60324764 1.83361363 1.18538666 1.19108021]",
-                  "[1.6177492141723633, 1.8144152164459229, 1.1718573570251465, 1.1950569152832031]"
+                  "[0.79905868 0.89367443 0.75429088 0.3190186 ]",
+                  "[1.79905868 1.89367437 1.75429082 1.3190186 ]",
+                  "[1.5757312774658203, 1.7105278968811035, 1.7636661529541016, 1.3394038677215576]"
                 ],
                 [
-                  "[0.15064228 0.03198934 0.25754827 0.51484001]",
-                  "[1.15064228 1.03198934 1.25754833 1.51484001]",
-                  "[1.1556042432785034, 0.9955940246582031, 1.2316606044769287, 1.5150485038757324]"
+                  "[0.80893755 0.92237449 0.88346356 0.93164903]",
+                  "[1.80893755 1.92237449 1.88346362 1.93164897]",
+                  "[1.562132716178894, 1.6031286716461182, 1.593322992324829, 1.4810831546783447]"
                 ],
                 [
-                  "[0.48959708 0.48549271 0.32688856 0.356677  ]",
-                  "[1.48959708 1.48549271 1.32688856 1.35667706]",
-                  "[1.4930214881896973, 1.467790961265564, 1.3132573366165161, 1.3589863777160645]"
+                  "[0.26661873 0.45946234 0.13510543 0.81294441]",
+                  "[1.26661873 1.4594624  1.13510537 1.81294441]",
+                  "[1.533058762550354, 1.3753284215927124, 1.230975866317749, 1.7815138101577759]"
                 ],
                 [
-                  "[0.08107251 0.2602725  0.18861133 0.44833237]",
-                  "[1.08107257 1.2602725  1.18861127 1.44833231]",
-                  "[1.102121114730835, 1.2180893421173096, 1.160165548324585, 1.4495322704315186]"
+                  "[0.39147133 0.29854035 0.84663737 0.58175623]",
+                  "[1.39147139 1.29854035 1.84663737 1.58175623]",
+                  "[1.5607244968414307, 1.5942375659942627, 1.5779708623886108, 1.4935153722763062]"
                 ],
                 [
-                  "[0.68094063 0.45189077 0.22661722 0.37354094]",
-                  "[1.68094063 1.45189071 1.22661722 1.37354088]",
-                  "[1.6725687980651855, 1.4393560886383057, 1.2169336080551147, 1.3746893405914307]"
+                  "[0.34084332 0.73018837 0.54168713 0.91440833]",
+                  "[1.34084332 1.73018837 1.54168713 1.91440833]",
+                  "[1.5488454103469849, 1.4963982105255127, 1.422922968864441, 1.622254490852356]"
                 ]
               ]
             },
@@ -688,13 +688,17 @@
                   "[0.11693293 0.49860179 0.55020827 0.88832849]",
                   "[1.11693287 1.49860179 1.55020833 1.88832855]"
                 ],
+                [
+                  "[0.48959708 0.48549271 0.32688856 0.356677  ]",
+                  "[1.48959708 1.48549271 1.32688856 1.35667706]"
+                ],
                 [
                   "[0.50272274 0.54912758 0.17663097 0.79070699]",
                   "[1.50272274 1.54912758 1.17663097 1.79070699]"
                 ],
                 [
-                  "[0.19908059 0.17570406 0.51475513 0.1893943 ]",
-                  "[1.19908059 1.175704   1.51475513 1.18939424]"
+                  "[0.04508126 0.76880038 0.80721325 0.62542385]",
+                  "[1.04508126 1.76880038 1.80721331 1.62542391]"
                 ],
                 [
                   "[0.40167677 0.25953674 0.9407078  0.76308483]",
@@ -712,10 +716,18 @@
                   "[0.62569475 0.9881897  0.83639616 0.9828859 ]",
                   "[1.62569475 1.9881897  1.83639622 1.98288584]"
                 ],
+                [
+                  "[0.56922203 0.98222166 0.76851749 0.28615737]",
+                  "[1.56922197 1.9822216  1.76851749 1.28615737]"
+                ],
                 [
                   "[0.88776821 0.51636773 0.30333066 0.32230979]",
                   "[1.88776827 1.51636767 1.30333066 1.32230973]"
                 ],
+                [
+                  "[0.90817457 0.89270043 0.38583666 0.66566533]",
+                  "[1.90817451 1.89270043 1.3858366  1.66566539]"
+                ],
                 [
                   "[0.48507756 0.80808765 0.77162558 0.47834778]",
                   "[1.48507762 1.80808759 1.77162552 1.47834778]"
@@ -724,10 +736,6 @@
                   "[0.68062544 0.98093534 0.14778823 0.53244978]",
                   "[1.68062544 1.98093534 1.14778829 1.53244972]"
                 ],
-                [
-                  "[0.31518555 0.49643308 0.11509258 0.95458382]",
-                  "[1.31518555 1.49643302 1.11509252 1.95458388]"
-                ],
                 [
                   "[0.79121011 0.54161114 0.69369799 0.1520769 ]",
                   "[1.79121017 1.54161119 1.69369793 1.15207696]"
@@ -744,10 +752,6 @@
                   "[0.94516498 0.08422136 0.5608117  0.07652664]",
                   "[1.94516492 1.08422136 1.56081176 1.07652664]"
                 ],
-                [
-                  "[0.26661873 0.45946234 0.13510543 0.81294441]",
-                  "[1.26661873 1.4594624  1.13510537 1.81294441]"
-                ],
                 [
                   "[0.30754459 0.77694583 0.09278506 0.38326019]",
                   "[1.30754459 1.77694583 1.09278512 1.38326025]"
@@ -804,10 +808,6 @@
                   "[0.73217702 0.65233225 0.44077861 0.33837909]",
                   "[1.73217702 1.65233231 1.44077861 1.33837914]"
                 ],
-                [
-                  "[0.34084332 0.73018837 0.54168713 0.91440833]",
-                  "[1.34084332 1.73018837 1.54168713 1.91440833]"
-                ],
                 [
                   "[0.60110539 0.3618983  0.32342511 0.98672163]",
                   "[1.60110545 1.3618983  1.32342505 1.98672163]"
@@ -816,6 +816,10 @@
                   "[0.77427191 0.21829212 0.12769502 0.74303615]",
                   "[1.77427197 1.21829212 1.12769508 1.74303615]"
                 ],
+                [
+                  "[0.08107251 0.2602725  0.18861133 0.44833237]",
+                  "[1.08107257 1.2602725  1.18861127 1.44833231]"
+                ],
                 [
                   "[0.59812403 0.78395379 0.0291847  0.81814629]",
                   "[1.59812403 1.78395379 1.0291847  1.81814623]"
@@ -840,18 +844,6 @@
                   "[0.95928186 0.84273899 0.71514636 0.38619852]",
                   "[1.95928192 1.84273899 1.7151463  1.38619852]"
                 ],
-                [
-                  "[0.32565445 0.90939188 0.07488042 0.13730896]",
-                  "[1.32565451 1.90939188 1.07488036 1.13730896]"
-                ],
-                [
-                  "[0.9829582  0.59269661 0.40120947 0.95487177]",
-                  "[1.9829582  1.59269667 1.40120947 1.95487177]"
-                ],
-                [
-                  "[0.79905868 0.89367443 0.75429088 0.3190186 ]",
-                  "[1.79905868 1.89367437 1.75429082 1.3190186 ]"
-                ],
                 [
                   "[0.54914117 0.03810108 0.87531954 0.73044223]",
                   "[1.54914117 1.03810108 1.87531948 1.73044229]"
@@ -876,10 +868,6 @@
                   "[0.60075855 0.12234765 0.00614399 0.30560958]",
                   "[1.60075855 1.12234759 1.00614405 1.30560958]"
                 ],
-                [
-                  "[0.39147133 0.29854035 0.84663737 0.58175623]",
-                  "[1.39147139 1.29854035 1.84663737 1.58175623]"
-                ],
                 [
                   "[0.02162331 0.81861657 0.92468154 0.07808572]",
                   "[1.02162337 1.81861663 1.92468154 1.07808566]"
@@ -924,6 +912,10 @@
                   "[0.59492421 0.90274489 0.38069052 0.46101224]",
                   "[1.59492421 1.90274489 1.38069057 1.46101224]"
                 ],
+                [
+                  "[0.15064228 0.03198934 0.25754827 0.51484001]",
+                  "[1.15064228 1.03198934 1.25754833 1.51484001]"
+                ],
                 [
                   "[0.12024075 0.21342516 0.56858408 0.58644271]",
                   "[1.12024069 1.21342516 1.56858408 1.58644271]"
@@ -932,6 +924,14 @@
                   "[0.91730917 0.22574073 0.09591609 0.33056474]",
                   "[1.91730917 1.22574067 1.09591603 1.33056474]"
                 ],
+                [
+                  "[0.49691743 0.61873293 0.90698647 0.94486356]",
+                  "[1.49691749 1.61873293 1.90698647 1.94486356]"
+                ],
+                [
+                  "[0.6032477  0.83361369 0.18538666 0.19108021]",
+                  "[1.60324764 1.83361363 1.18538666 1.19108021]"
+                ],
                 [
                   "[0.63235509 0.70352674 0.96188956 0.46240485]",
                   "[1.63235509 1.70352674 1.96188951 1.46240485]"
@@ -948,10 +948,6 @@
                   "[0.40234613 0.54987347 0.49542785 0.54153186]",
                   "[1.40234613 1.54987347 1.49542785 1.5415318 ]"
                 ],
-                [
-                  "[0.80893755 0.92237449 0.88346356 0.93164903]",
-                  "[1.80893755 1.92237449 1.88346362 1.93164897]"
-                ],
                 [
                   "[0.12858278 0.09930819 0.83222693 0.72485673]",
                   "[1.12858272 1.09930825 1.83222699 1.72485673]"
@@ -981,13 +977,17 @@
                   "[1.28942847 1.05601001 1.33039129 1.27781558]"
                 ],
                 [
-                  "[0.43681622 0.74680805 0.83598751 0.12414402]",
-                  "[1.43681622 1.74680805 1.83598757 1.12414408]"
+                  "[0.68094063 0.45189077 0.22661722 0.37354094]",
+                  "[1.68094063 1.45189071 1.22661722 1.37354088]"
                 ],
                 [
                   "[0.47870928 0.17129105 0.27300501 0.20634609]",
                   "[1.47870922 1.17129111 1.27300501 1.20634604]"
                 ],
+                [
+                  "[0.72795159 0.79317838 0.27832931 0.96576637]",
+                  "[1.72795153 1.79317832 1.27832937 1.96576643]"
+                ],
                 [
                   "[0.87608397 0.93200487 0.80169648 0.37758952]",
                   "[1.87608397 1.93200493 1.80169654 1.37758946]"
@@ -1000,7 +1000,7 @@
             }
           },
           "other": {
-            "model": "ModelConfig(model=Sequential(\n  (0) - Linear(in_features=4, out_features=4, bias=True): Input__embedding_1_x -> Linear_1_x\n  (1) - <function leaky_relu at 0x719e0ce23a60>: Linear_1_x -> Activation_2_x\n  (2) - Identity(): Activation_2_x -> Activation_2_x\n), model_inputs=['Input__embedding_1_x'], model_outputs=['Activation_2_x'], loss_inputs=['Input__label_1_y', 'Activation_2_x'], loss=Sequential(\n  (0) - <function mse_loss at 0x719e0ce2d580>: Activation_2_x, Input__label_1_y -> MSE_loss_1_loss\n  (1) - Identity(): MSE_loss_1_loss -> loss\n), optimizer=SGD (\nParameter Group 0\n    dampening: 0\n    differentiable: False\n    foreach: None\n    fused: None\n    lr: 0.1\n    maximize: False\n    momentum: 0\n    nesterov: False\n    weight_decay: 0\n), source_workspace=None, trained=True)"
+            "model": "ModelConfig(model=Sequential(\n  (0) - Identity(): Input__tensor_1_x -> START_Repeat_1_output\n  (1) - Linear(in_features=4, out_features=4, bias=True): START_Repeat_1_output -> Linear_1_output\n  (2) - <function leaky_relu at 0x762d1f82c680>: Linear_1_output -> Activation_1_output\n  (3) - Identity(): Activation_1_output -> START_Repeat_1_output\n  (4) - Linear(in_features=4, out_features=4, bias=True): START_Repeat_1_output -> Linear_1_output\n  (5) - <function leaky_relu at 0x762d1f82c680>: Linear_1_output -> Activation_1_output\n  (6) - Identity(): Activation_1_output -> END_Repeat_1_output\n  (7) - Identity(): END_Repeat_1_output -> END_Repeat_1_output\n), model_inputs=['Input__tensor_1_x'], model_outputs=['END_Repeat_1_output'], loss_inputs=['END_Repeat_1_output', 'Input__tensor_3_x'], loss=Sequential(\n  (0) - <function mse_loss at 0x762d1f82e160>: END_Repeat_1_output, Input__tensor_3_x -> MSE_loss_2_output\n  (1) - Identity(): MSE_loss_2_output -> loss\n), optimizer_parameters={'lr': 0.1, 'type': <OptionsFor_type.SGD: 4>}, optimizer=SGD (\nParameter Group 0\n    dampening: 0\n    differentiable: False\n    foreach: None\n    fused: None\n    lr: 0.1\n    maximize: False\n    momentum: 0\n    nesterov: False\n    weight_decay: 0\n), source_workspace='Model definition', trained=True)"
           },
           "relations": []
         },
@@ -1016,7 +1016,7 @@
               },
               "df_test": {
                 "columns": [
-                  "predicted",
+                  "pred",
                   "x",
                   "y"
                 ]
@@ -1032,14 +1032,14 @@
               "model": {
                 "model": {
                   "inputs": [
-                    "Input__embedding_1_x"
+                    "Input__tensor_1_x"
                   ],
                   "loss_inputs": [
-                    "Input__label_1_y",
-                    "Activation_2_x"
+                    "END_Repeat_1_output",
+                    "Input__tensor_3_x"
                   ],
                   "outputs": [
-                    "Activation_2_x"
+                    "END_Repeat_1_output"
                   ],
                   "trained": true
                 },
@@ -1207,14 +1207,14 @@
               "model": {
                 "model": {
                   "inputs": [
-                    "Input__embedding_1_x"
+                    "Input__tensor_1_x"
                   ],
                   "loss_inputs": [
-                    "Input__label_1_y",
-                    "Activation_2_x"
+                    "END_Repeat_1_output",
+                    "Input__tensor_3_x"
                   ],
                   "outputs": [
-                    "Activation_2_x"
+                    "END_Repeat_1_output"
                   ],
                   "trained": false
                 },
@@ -1270,8 +1270,8 @@
           "type": "basic"
         },
         "params": {
-          "epochs": "1001",
-          "input_mapping": "{\"map\":{\"Input__embedding_1_x\":{\"column\":\"x\",\"df\":\"df_train\"},\"Input__label_1_y\":{\"column\":\"y\",\"df\":\"df_train\"}}}",
+          "epochs": "1003",
+          "input_mapping": "{\"map\":{\"Input__tensor_1_x\":{\"df\":\"df_train\",\"column\":\"x\"},\"Input__tensor_3_x\":{\"df\":\"df_train\",\"column\":\"y\"}}}",
           "model_name": "model"
         },
         "status": "done",
@@ -1304,7 +1304,6 @@
               },
               "df_test": {
                 "columns": [
-                  "predicted",
                   "x",
                   "y"
                 ]
@@ -1320,14 +1319,14 @@
               "model": {
                 "model": {
                   "inputs": [
-                    "Input__embedding_1_x"
+                    "Input__tensor_1_x"
                   ],
                   "loss_inputs": [
-                    "Input__label_1_y",
-                    "Activation_2_x"
+                    "END_Repeat_1_output",
+                    "Input__tensor_3_x"
                   ],
                   "outputs": [
-                    "Activation_2_x"
+                    "END_Repeat_1_output"
                   ],
                   "trained": true
                 },
@@ -1383,15 +1382,15 @@
           "type": "basic"
         },
         "params": {
-          "input_mapping": "{\"map\":{\"Input__embedding_1_x\":{\"column\":\"x\",\"df\":\"df_test\"}}}",
+          "input_mapping": "{\"map\":{\"Input__tensor_1_x\":{\"df\":\"df_test\",\"column\":\"x\"}}}",
           "model_name": "model",
-          "output_mapping": "{\"map\":{\"Activation_2_x\":{\"column\":\"predicted\",\"df\":\"df_test\"}}}"
+          "output_mapping": "{\"map\":{\"END_Repeat_1_output\":{\"df\":\"df_test\",\"column\":\"pred\"}}}"
         },
         "status": "done",
         "title": "Model inference"
       },
       "dragHandle": ".bg-primary",
-      "height": 893.0,
+      "height": 650.0,
       "id": "Model inference 1",
       "position": {
         "x": 2181.718373860645,

lynxkite-app/web/src/workspace/nodes/NodeParameter.tsx

@@ -1,3 +1,4 @@
+import { useRef } from "react";
 // @ts-ignore
 import ArrowsHorizontal from "~icons/tabler/arrows-horizontal.jsx";
 
@@ -14,13 +15,16 @@ function ParamName({ name }: { name: string }) {
 function Input({
   value,
   onChange,
+  inputRef,
 }: {
   value: string;
   onChange: (value: string, options?: { delay: number }) => void;
+  inputRef?: React.Ref<HTMLInputElement>;
 }) {
   return (
     <input
       className="input input-bordered w-full"
+      ref={inputRef}
       value={value || ""}
       onChange={(evt) => onChange(evt.currentTarget.value, { delay: 2 })}
       onBlur={(evt) => onChange(evt.currentTarget.value, { delay: 0 })}
@@ -29,6 +33,13 @@ function Input({
   );
 }
 
+type Bindings = {
+  [key: string]: {
+    df: string;
+    column: string;
+  };
+};
+
 function getModelBindings(
   data: any,
   variant: "training input" | "inference input" | "output",
@@ -71,11 +82,16 @@ function parseJsonOrEmpty(json: string): object {
 }
 
 function ModelMapping({ value, onChange, data, variant }: any) {
+  const dfsRef = useRef({} as { [binding: string]: HTMLSelectElement | null });
+  const columnsRef = useRef(
+    {} as { [binding: string]: HTMLSelectElement | HTMLInputElement | null },
+  );
   const v: any = parseJsonOrEmpty(value);
   v.map ??= {};
   const dfs: { [df: string]: string[] } = {};
   const inputs = data?.input_metadata?.value ?? data?.input_metadata ?? [];
   for (const input of inputs) {
+    if (!input.dataframes) continue;
     const dataframes = input.dataframes as {
       [df: string]: { columns: string[] };
     };
@@ -84,6 +100,17 @@ function ModelMapping({ value, onChange, data, variant }: any) {
     }
   }
   const bindings = getModelBindings(data, variant);
+  function getMap() {
+    const map: Bindings = {};
+    for (const binding of bindings) {
+      const df = dfsRef.current[binding]?.value ?? "";
+      const column = columnsRef.current[binding]?.value ?? "";
+      if (df.length || column.length) {
+        map[binding] = { df, column };
+      }
+    }
+    return map;
+  }
   return (
     <table className="model-mapping-param">
       <tbody>
@@ -98,21 +125,10 @@ function ModelMapping({ value, onChange, data, variant }: any) {
                 <select
                   className="select select-ghost"
                   value={v.map?.[binding]?.df}
-                  onChange={(evt) => {
-                    const df = evt.currentTarget.value;
-                    if (df === "") {
-                      const map = { ...v.map, [binding]: undefined };
-                      onChange(JSON.stringify({ map }));
-                    } else {
-                      const columnSpec = {
-                        column: dfs[df][0],
-                        ...(v.map?.[binding] || {}),
-                        df,
-                      };
-                      const map = { ...v.map, [binding]: columnSpec };
-                      onChange(JSON.stringify({ map }));
-                    }
+                  ref={(el) => {
+                    dfsRef.current[binding] = el;
                   }}
+                  onChange={() => onChange(JSON.stringify({ map: getMap() }))}
                 >
                   <option key="" value="" />
                   {Object.keys(dfs).map((df: string) => (
@@ -125,13 +141,16 @@ function ModelMapping({ value, onChange, data, variant }: any) {
               <td>
                 {variant === "output" ? (
                   <Input
+                    inputRef={(el) => {
+                      columnsRef.current[binding] = el;
+                    }}
                     value={v.map?.[binding]?.column}
                     onChange={(column, options) => {
-                      const columnSpec = {
-                        ...(v.map?.[binding] || {}),
-                        column,
-                      };
-                      const map = { ...v.map, [binding]: columnSpec };
+                      const map = getMap();
+                      // At this point the <input> has not been updated yet. We use the value from the event.
+                      const df = dfsRef.current[binding]?.value ?? "";
+                      map[binding] ??= { df, column };
+                      map[binding].column = column;
                       onChange(JSON.stringify({ map }), options);
                     }}
                   />
@@ -139,16 +158,12 @@ function ModelMapping({ value, onChange, data, variant }: any) {
                   <select
                     className="select select-ghost"
                     value={v.map?.[binding]?.column}
-                    onChange={(evt) => {
-                      const column = evt.currentTarget.value;
-                      const columnSpec = {
-                        ...(v.map?.[binding] || {}),
-                        column,
-                      };
-                      const map = { ...v.map, [binding]: columnSpec };
-                      onChange(JSON.stringify({ map }));
+                    ref={(el) => {
+                      columnsRef.current[binding] = el;
                     }}
+                    onChange={() => onChange(JSON.stringify({ map: getMap() }))}
                   >
+                    <option key="" value="" />
                     {dfs[v.map?.[binding]?.df]?.map((col: string) => (
                       <option key={col} value={col}>
                         {col}

lynxkite-core/src/lynxkite/core/ops.py

@@ -1,6 +1,7 @@
 """API for implementing LynxKite operations."""
 
 from __future__ import annotations
+import asyncio
 import enum
 import functools
 import inspect
@@ -13,7 +14,7 @@ from typing_extensions import Annotated
 if typing.TYPE_CHECKING:
     from . import workspace
 
-CATALOGS = {}
+CATALOGS: dict[str, dict[str, "Op"]] = {}
 EXECUTORS = {}
 
 typeof = type  # We have some arguments called "type".
@@ -297,3 +298,13 @@ def op_registration(env: str):
 def passive_op_registration(env: str):
     """Returns a function that can be used to register operations without associated code."""
     return functools.partial(register_passive_op, env)
+
+
+def slow(func):
+    """Decorator for slow, blocking operations. Turns them into separate threads."""
+
+    @functools.wraps(func)
+    async def wrapper(*args, **kwargs):
+        return await asyncio.to_thread(func, *args, **kwargs)
+
+    return wrapper

lynxkite-graph-analytics/src/lynxkite_graph_analytics/core.py

@@ -1,5 +1,6 @@
 """Graph analytics executor and data types."""
 
+import inspect
 import os
 from lynxkite.core import ops, workspace
 import dataclasses
@@ -177,10 +178,16 @@ async def execute(ws: workspace.Workspace):
                 # All inputs for this node are ready, we can compute the output.
                 todo.remove(id)
                 progress = True
-                _execute_node(node, ws, catalog, outputs)
+                await _execute_node(node, ws, catalog, outputs)
 
 
-def _execute_node(node, ws, catalog, outputs):
+async def await_if_needed(obj):
+    if inspect.isawaitable(obj):
+        obj = await obj
+    return obj
+
+
+async def _execute_node(node, ws, catalog, outputs):
     params = {**node.data.params}
     op = catalog.get(node.data.title)
     if not op:
@@ -214,6 +221,7 @@ def _execute_node(node, ws, catalog, outputs):
     # Execute op.
     try:
         result = op(*inputs, **params)
+        result.output = await await_if_needed(result.output)
     except Exception as e:
         if os.environ.get("LYNXKITE_LOG_OP_ERRORS"):
             traceback.print_exc()

lynxkite-graph-analytics/src/lynxkite_graph_analytics/lynxkite_ops.py

@@ -347,7 +347,7 @@ def define_model(
     assert model_workspace, "Model workspace is unset."
     ws = load_ws(model_workspace)
     # Build the model without inputs, to get its interface.
-    m = pytorch_model_ops.build_model(ws, {})
+    m = pytorch_model_ops.build_model(ws)
     m.source_workspace = model_workspace
     bundle = bundle.copy()
     bundle.other[save_as] = m
@@ -369,6 +369,7 @@ class ModelOutputMapping(pytorch_model_ops.ModelMapping):
 
 
 @op("Train model")
+@ops.slow
 def train_model(
     bundle: core.Bundle,
     *,
@@ -379,14 +380,12 @@ def train_model(
     """Trains the selected model on the selected dataset. Most training parameters are set in the model definition."""
     m = bundle.other[model_name].copy()
     inputs = pytorch_model_ops.to_tensors(bundle, input_mapping)
-    if not m.trained and m.source_workspace:
-        # Rebuild the model for the correct inputs.
-        ws = load_ws(m.source_workspace)
-        m = pytorch_model_ops.build_model(ws, inputs)
     t = tqdm(range(epochs), desc="Training model")
+    losses = []
     for _ in t:
         loss = m.train(inputs)
         t.set_postfix({"loss": loss})
+        losses.append(loss)
     m.trained = True
     bundle = bundle.copy()
     bundle.other[model_name] = m
@@ -394,6 +393,7 @@ def train_model(
 
 
 @op("Model inference")
+@ops.slow
 def model_inference(
     bundle: core.Bundle,
     *,
@@ -409,7 +409,13 @@ def model_inference(
     inputs = pytorch_model_ops.to_tensors(bundle, input_mapping)
     outputs = m.inference(inputs)
     bundle = bundle.copy()
+    copied = set()
     for k, v in output_mapping.map.items():
+        if not v.df or not v.column:
+            continue
+        if v.df not in copied:
+            bundle.dfs[v.df] = bundle.dfs[v.df].copy()
+            copied.add(v.df)
         bundle.dfs[v.df][v.column] = outputs[k].detach().numpy().tolist()
     return bundle
 

lynxkite-graph-analytics/src/lynxkite_graph_analytics/pytorch_model_ops.py

@@ -1,20 +1,35 @@
 """Boxes for defining PyTorch models."""
 
 import copy
+import enum
 import graphlib
-import types
 
 import pydantic
 from lynxkite.core import ops, workspace
 from lynxkite.core.ops import Parameter as P
 import torch
-import torch_geometric as pyg
+import torch_geometric.nn as pyg_nn
 import dataclasses
 from . import core
 
 ENV = "PyTorch model"
 
 
+def op(name, **kwargs):
+    _op = ops.op(ENV, name, **kwargs)
+
+    def decorator(func):
+        _op(func)
+        op = func.__op__
+        for p in op.inputs.values():
+            p.position = "bottom"
+        for p in op.outputs.values():
+            p.position = "top"
+        return func
+
+    return decorator
+
+
 def reg(name, inputs=[], outputs=None, params=[]):
     if outputs is None:
         outputs = inputs
@@ -27,13 +42,9 @@ def reg(name, inputs=[], outputs=None, params=[]):
     )
 
 
-reg("Input: embedding", outputs=["x"])
+reg("Input: tensor", outputs=["output"], params=[P.basic("name")])
 reg("Input: graph edges", outputs=["edges"])
-reg("Input: label", outputs=["y"])
-reg("Input: positive sample", outputs=["x_pos"])
-reg("Input: negative sample", outputs=["x_neg"])
 reg("Input: sequential", outputs=["y"])
-reg("Input: zeros", outputs=["x"])
 
 reg("LSTM", inputs=["x", "h"], outputs=["x", "h"])
 reg(
@@ -59,10 +70,35 @@ reg(
         ),
     ],
 )
+
+
 reg("Attention", inputs=["q", "k", "v"], outputs=["x", "weights"])
 reg("LayerNorm", inputs=["x"])
 reg("Dropout", inputs=["x"], params=[P.basic("p", 0.5)])
-reg("Linear", inputs=["x"], params=[P.basic("output_dim", "same")])
+
+
+@op("Linear")
+def linear(x, *, output_dim=1024):
+    return pyg_nn.Linear(-1, output_dim)
+
+
+class ActivationTypes(enum.Enum):
+    ReLU = "ReLU"
+    Leaky_ReLU = "Leaky ReLU"
+    Tanh = "Tanh"
+    Mish = "Mish"
+
+
+@op("Activation")
+def activation(x, *, type: ActivationTypes = ActivationTypes.ReLU):
+    return getattr(torch.nn.functional, type.name.lower().replace(" ", "_"))
+
+
+@op("MSE loss")
+def mse_loss(x, y):
+    return torch.nn.functional.mse_loss
+
+
 reg("Softmax", inputs=["x"])
 reg(
     "Graph conv",
@@ -70,16 +106,10 @@ reg(
     outputs=["x"],
     params=[P.options("type", ["GCNConv", "GATConv", "GATv2Conv", "SAGEConv"])],
 )
-reg(
-    "Activation",
-    inputs=["x"],
-    params=[P.options("type", ["ReLU", "Leaky ReLU", "Tanh", "Mish"])],
-)
 reg("Concatenate", inputs=["a", "b"], outputs=["x"])
 reg("Add", inputs=["a", "b"], outputs=["x"])
 reg("Subtract", inputs=["a", "b"], outputs=["x"])
 reg("Multiply", inputs=["a", "b"], outputs=["x"])
-reg("MSE loss", inputs=["x", "y"], outputs=["loss"])
 reg("Triplet margin loss", inputs=["x", "x_pos", "x_neg"], outputs=["loss"])
 reg("Cross-entropy loss", inputs=["x", "y"], outputs=["loss"])
 reg(
@@ -110,7 +140,7 @@ ops.register_passive_op(
     outputs=[ops.Output(name="output", position="bottom", type="tensor")],
     params=[
         ops.Parameter.basic("times", 1, int),
-        ops.Parameter.basic("same_weights", True, bool),
+        ops.Parameter.basic("same_weights", False, bool),
     ],
 )
 
@@ -128,6 +158,21 @@ def _to_id(*strings: str) -> str:
     return "_".join("".join(c if c.isalnum() else "_" for c in s) for s in strings)
 
 
+@dataclasses.dataclass
+class Layer:
+    """Temporary data structure used by ModelBuilder."""
+
+    module: torch.nn.Module
+    origin_id: str
+    inputs: list[str]
+    outputs: list[str]
+
+    def for_sequential(self):
+        inputs = ", ".join(self.inputs)
+        outputs = ", ".join(self.outputs)
+        return self.module, f"{inputs} -> {outputs}"
+
+
 class ColumnSpec(pydantic.BaseModel):
     df: str
     column: str
@@ -144,10 +189,17 @@ class ModelConfig:
     model_outputs: list[str]
     loss_inputs: list[str]
     loss: torch.nn.Module
-    optimizer: torch.optim.Optimizer
+    optimizer_parameters: dict[str, any]
+    optimizer: torch.optim.Optimizer | None = None
     source_workspace: str | None = None
     trained: bool = False
 
+    def __post_init__(self):
+        self._make_optimizer()
+
+    def num_parameters(self) -> int:
+        return sum(p.numel() for p in self.model.parameters())
+
     def _forward(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
         model_inputs = [inputs[i] for i in self.model_inputs]
         output = self.model(*model_inputs)
@@ -174,10 +226,20 @@ class ModelConfig:
         self.optimizer.step()
         return loss.item()
 
+    def _make_optimizer(self):
+        # We need to make a new optimizer when the model is copied. (It's tied to its parameters.)
+        p = self.optimizer_parameters
+        o = getattr(torch.optim, p["type"].name)
+        self.optimizer = o(self.model.parameters(), lr=p["lr"])
+
     def copy(self):
         """Returns a copy of the model."""
-        c = dataclasses.replace(self)
-        c.model = copy.deepcopy(self.model)
+        c = dataclasses.replace(
+            self,
+            model=copy.deepcopy(self.model),
+        )
+        c._make_optimizer()
+        c.optimizer.load_state_dict(self.optimizer.state_dict())
         return c
 
     def metadata(self):
@@ -192,113 +254,219 @@ class ModelConfig:
         }
 
 
-def build_model(ws: workspace.Workspace, inputs: dict[str, torch.Tensor]) -> ModelConfig:
+def build_model(ws: workspace.Workspace) -> ModelConfig:
     """Builds the model described in the workspace."""
-    catalog = ops.CATALOGS[ENV]
-    optimizers = []
-    nodes = {}
-    for node in ws.nodes:
-        nodes[node.id] = node
-        if node.data.title == "Optimizer":
-            optimizers.append(node.id)
-    assert optimizers, "No optimizer found."
-    assert len(optimizers) == 1, f"More than one optimizer found: {optimizers}"
-    [optimizer] = optimizers
-    dependencies = {n.id: [] for n in ws.nodes}
-    in_edges = {}
-    out_edges = {}
-    # TODO: Dissolve repeat boxes here.
-    for e in ws.edges:
-        dependencies[e.target].append(e.source)
-        in_edges.setdefault(e.target, {}).setdefault(e.targetHandle, []).append(
-            (e.source, e.sourceHandle)
-        )
-        out_edges.setdefault(e.source, {}).setdefault(e.sourceHandle, []).append(
-            (e.target, e.targetHandle)
-        )
-    sizes = {}
-    for k, i in inputs.items():
-        sizes[k] = i.shape[-1]
-    ts = graphlib.TopologicalSorter(dependencies)
-    layers = []
-    loss_layers = []
-    in_loss = set()
-    cfg = {}
-    used_in_model = set()
-    made_in_model = set()
-    used_in_loss = set()
-    made_in_loss = set()
-    for node_id in ts.static_order():
-        node = nodes[node_id]
+    builder = ModelBuilder(ws)
+    return builder.build_model()
+
+
+class ModelBuilder:
+    """The state shared between methods that are used to build the model."""
+
+    def __init__(self, ws: workspace.Workspace):
+        self.catalog = ops.CATALOGS[ENV]
+        optimizers = []
+        self.nodes: dict[str, workspace.WorkspaceNode] = {}
+        repeats: list[str] = []
+        for node in ws.nodes:
+            self.nodes[node.id] = node
+            if node.data.title == "Optimizer":
+                optimizers.append(node.id)
+            elif node.data.title == "Repeat":
+                repeats.append(node.id)
+                self.nodes[f"START {node.id}"] = node
+                self.nodes[f"END {node.id}"] = node
+        assert optimizers, "No optimizer found."
+        assert len(optimizers) == 1, f"More than one optimizer found: {optimizers}"
+        [self.optimizer] = optimizers
+        self.dependencies = {n: [] for n in self.nodes}
+        self.in_edges: dict[str, dict[str, list[tuple[str, str]]]] = {n: {} for n in self.nodes}
+        self.out_edges: dict[str, dict[str, list[tuple[str, str]]]] = {n: {} for n in self.nodes}
+        for e in ws.edges:
+            self.dependencies[e.target].append(e.source)
+            self.in_edges.setdefault(e.target, {}).setdefault(e.targetHandle, []).append(
+                (e.source, e.sourceHandle)
+            )
+            self.out_edges.setdefault(e.source, {}).setdefault(e.sourceHandle, []).append(
+                (e.target, e.targetHandle)
+            )
+        # Split repeat boxes into start and end, and insert them into the flow.
+        # TODO: Think about recursive repeats.
+        for repeat in repeats:
+            if not self.out_edges[repeat] or not self.in_edges[repeat]:
+                continue
+            start_id = f"START {repeat}"
+            end_id = f"END {repeat}"
+            # repeat -> first <- real_input
+            # ...becomes...
+            # real_input -> start -> first
+            first, firsth = self.out_edges[repeat]["output"][0]
+            [(real_input, real_inputh)] = [
+                k for k in self.in_edges[first][firsth] if k != (repeat, "output")
+            ]
+            self.dependencies[first].remove(repeat)
+            self.dependencies[first].append(start_id)
+            self.dependencies[start_id] = [real_input]
+            self.out_edges[real_input][real_inputh] = [
+                k if k != (first, firsth) else (start_id, "input")
+                for k in self.out_edges[real_input][real_inputh]
+            ]
+            self.in_edges[start_id] = {"input": [(real_input, real_inputh)]}
+            self.out_edges[start_id] = {"output": [(first, firsth)]}
+            self.in_edges[first][firsth] = [(start_id, "output")]
+            # repeat <- last -> real_output
+            # ...becomes...
+            # last -> end -> real_output
+            last, lasth = self.in_edges[repeat]["input"][0]
+            [(real_output, real_outputh)] = [
+                k for k in self.out_edges[last][lasth] if k != (repeat, "input")
+            ]
+            del self.dependencies[repeat]
+            self.dependencies[end_id] = [last]
+            self.dependencies[real_output].append(end_id)
+            self.out_edges[last][lasth] = [(end_id, "input")]
+            self.in_edges[end_id] = {"input": [(last, lasth)]}
+            self.out_edges[end_id] = {"output": [(real_output, real_outputh)]}
+            self.in_edges[real_output][real_outputh] = [
+                k if k != (last, lasth) else (end_id, "output")
+                for k in self.in_edges[real_output][real_outputh]
+            ]
+        self.inv_dependencies = {n: [] for n in self.nodes}
+        for k, v in self.dependencies.items():
+            for i in v:
+                self.inv_dependencies[i].append(k)
+        self.layers: list[Layer] = []
+        # Clean up disconnected nodes.
+        disconnected = set()
+        for node_id in self.nodes:
+            op = self.catalog[self.nodes[node_id].data.title]
+            if len(self.in_edges[node_id]) != len(op.inputs):
+                disconnected.add(node_id)
+                disconnected |= self.all_upstream(node_id)
+        for node_id in disconnected:
+            del self.dependencies[node_id]
+            del self.in_edges[node_id]
+            del self.out_edges[node_id]
+            del self.inv_dependencies[node_id]
+            del self.nodes[node_id]
+
+    def all_upstream(self, node: str) -> set[str]:
+        """Returns all nodes upstream of a node."""
+        deps = set()
+        for dep in self.dependencies[node]:
+            deps.add(dep)
+            deps.update(self.all_upstream(dep))
+        return deps
+
+    def all_downstream(self, node: str) -> set[str]:
+        """Returns all nodes downstream of a node."""
+        deps = set()
+        for dep in self.inv_dependencies[node]:
+            deps.add(dep)
+            deps.update(self.all_downstream(dep))
+        return deps
+
+    def run_node(self, node_id: str) -> None:
+        """Adds the layer(s) produced by this node to self.layers."""
+        node = self.nodes[node_id]
         t = node.data.title
-        op = catalog[t]
+        op = self.catalog[t]
         p = op.convert_params(node.data.params)
-        for b in dependencies[node_id]:
-            if b in in_loss:
-                in_loss.add(node_id)
-        if "loss" in t:
-            in_loss.add(node_id)
-        inputs = {}
-        for n in in_edges.get(node_id, []):
-            for b, h in in_edges[node_id][n]:
-                i = _to_id(b, h)
-                inputs[n] = i
-                if node_id in in_loss:
-                    used_in_loss.add(i)
-                else:
-                    used_in_model.add(i)
-        outputs = {}
-        for out in out_edges.get(node_id, []):
-            i = _to_id(node_id, out)
-            outputs[out] = i
-            if inputs:  # Nodes with no inputs are input nodes. Their outputs are not "made" by us.
-                if node_id in in_loss:
-                    made_in_loss.add(i)
-                else:
-                    made_in_model.add(i)
-        inputs = types.SimpleNamespace(**inputs)
-        outputs = types.SimpleNamespace(**outputs)
-        ls = loss_layers if node_id in in_loss else layers
         match t:
-            case "Linear":
-                isize = sizes.get(inputs.x, 1)
-                osize = isize if p["output_dim"] == "same" else int(p["output_dim"])
-                ls.append((torch.nn.Linear(isize, osize), f"{inputs.x} -> {outputs.x}"))
-                sizes[outputs.x] = osize
-            case "Activation":
-                f = getattr(torch.nn.functional, p["type"].name.lower().replace(" ", "_"))
-                ls.append((f, f"{inputs.x} -> {outputs.x}"))
-                sizes[outputs.x] = sizes.get(inputs.x, 1)
-            case "MSE loss":
-                ls.append(
-                    (
-                        torch.nn.functional.mse_loss,
-                        f"{inputs.x}, {inputs.y} -> {outputs.loss}",
+            case "Repeat":
+                if node_id.startswith("END "):
+                    repeat_id = node_id.removeprefix("END ")
+                    start_id = f"START {repeat_id}"
+                    [last_output] = self.in_edges[node_id]["input"]
+                    after_start = self.all_downstream(start_id)
+                    after_end = self.all_downstream(node_id)
+                    before_end = self.all_upstream(node_id)
+                    affected_nodes = after_start - after_end - {node_id}
+                    repeated_nodes = after_start & before_end
+                    assert affected_nodes == repeated_nodes, (
+                        f"edges leave repeated section '{repeat_id}':\n{affected_nodes - repeated_nodes}"
                     )
-                )
-    cfg["model_inputs"] = list(used_in_model - made_in_model)
-    cfg["model_outputs"] = list(made_in_model & used_in_loss)
-    cfg["loss_inputs"] = list(used_in_loss - made_in_loss)
-    # Make sure the trained output is output from the last model layer.
-    outputs = ", ".join(cfg["model_outputs"])
-    layers.append((torch.nn.Identity(), f"{outputs} -> {outputs}"))
-    # Create model.
-    cfg["model"] = pyg.nn.Sequential(", ".join(cfg["model_inputs"]), layers)
-    # Make sure the loss is output from the last loss layer.
-    [(lossb, lossh)] = in_edges[optimizer]["loss"]
-    lossi = _to_id(lossb, lossh)
-    loss_layers.append((torch.nn.Identity(), f"{lossi} -> loss"))
-    # Create loss function.
-    cfg["loss"] = pyg.nn.Sequential(", ".join(cfg["loss_inputs"]), loss_layers)
-    assert not list(cfg["loss"].parameters()), (
-        f"loss should have no parameters: {list(cfg['loss'].parameters())}"
-    )
-    # Create optimizer.
-    op = catalog["Optimizer"]
-    p = op.convert_params(nodes[optimizer].data.params)
-    o = getattr(torch.optim, p["type"].name)
-    cfg["optimizer"] = o(cfg["model"].parameters(), lr=p["lr"])
-    return ModelConfig(**cfg)
+                    repeated_layers = [e for e in self.layers if e.origin_id in repeated_nodes]
+                    assert p["times"] >= 1, f"Cannot repeat {repeat_id} {p['times']} times."
+                    for i in range(p["times"] - 1):
+                        # Copy repeat section's output to repeat section's input.
+                        self.layers.append(
+                            Layer(
+                                torch.nn.Identity(),
+                                origin_id=node_id,
+                                inputs=[_to_id(*last_output)],
+                                outputs=[_to_id(start_id, "output")],
+                            )
+                        )
+                        # Repeat the layers in the section.
+                        for layer in repeated_layers:
+                            if p["same_weights"]:
+                                self.layers.append(layer)
+                            else:
+                                self.run_node(layer.origin_id)
+                self.layers.append(self.run_op(node_id, op, p))
+            case "Optimizer" | "Input: tensor" | "Input: graph edges" | "Input: sequential":
+                return
+            case _:
+                self.layers.append(self.run_op(node_id, op, p))
+
+    def run_op(self, node_id: str, op: ops.Op, params) -> Layer:
+        """Returns the layer produced by this op."""
+        inputs = [_to_id(*i) for n in op.inputs for i in self.in_edges[node_id][n]]
+        outputs = [_to_id(node_id, n) for n in op.outputs]
+        if op.func == ops.no_op:
+            module = torch.nn.Identity()
+        else:
+            module = op.func(*inputs, **params)
+        return Layer(module, node_id, inputs, outputs)
+
+    def build_model(self) -> ModelConfig:
+        # Walk the graph in topological order.
+        ts = graphlib.TopologicalSorter(self.dependencies)
+        for node_id in ts.static_order():
+            self.run_node(node_id)
+        return self.get_config()
+
+    def get_config(self) -> ModelConfig:
+        # Split the design into model and loss.
+        loss_nodes = set()
+        for node_id in self.nodes:
+            if "loss" in self.nodes[node_id].data.title:
+                loss_nodes.add(node_id)
+                loss_nodes |= self.all_downstream(node_id)
+        layers = []
+        loss_layers = []
+        for layer in self.layers:
+            if layer.origin_id in loss_nodes:
+                loss_layers.append(layer)
+            else:
+                layers.append(layer)
+        used_in_model = set(input for layer in layers for input in layer.inputs)
+        used_in_loss = set(input for layer in loss_layers for input in layer.inputs)
+        made_in_model = set(output for layer in layers for output in layer.outputs)
+        made_in_loss = set(output for layer in loss_layers for output in layer.outputs)
+        layers = [layer.for_sequential() for layer in layers]
+        loss_layers = [layer.for_sequential() for layer in loss_layers]
+        cfg = {}
+        cfg["model_inputs"] = list(used_in_model - made_in_model)
+        cfg["model_outputs"] = list(made_in_model & used_in_loss)
+        cfg["loss_inputs"] = list(used_in_loss - made_in_loss)
+        # Make sure the trained output is output from the last model layer.
+        outputs = ", ".join(cfg["model_outputs"])
+        layers.append((torch.nn.Identity(), f"{outputs} -> {outputs}"))
+        # Create model.
+        cfg["model"] = pyg_nn.Sequential(", ".join(cfg["model_inputs"]), layers)
+        # Make sure the loss is output from the last loss layer.
+        [(lossb, lossh)] = self.in_edges[self.optimizer]["loss"]
+        lossi = _to_id(lossb, lossh)
+        loss_layers.append((torch.nn.Identity(), f"{lossi} -> loss"))
+        # Create loss function.
+        cfg["loss"] = pyg_nn.Sequential(", ".join(cfg["loss_inputs"]), loss_layers)
+        assert not list(cfg["loss"].parameters()), f"loss should have no parameters: {loss_layers}"
+        # Create optimizer.
+        op = self.catalog["Optimizer"]
+        cfg["optimizer_parameters"] = op.convert_params(self.nodes[self.optimizer].data.params)
+        return ModelConfig(**cfg)
 
 
 def to_tensors(b: core.Bundle, m: ModelMapping | None) -> dict[str, torch.Tensor]:

lynxkite-graph-analytics/tests/test_pytorch_model_ops.py

@@ -4,14 +4,16 @@ import torch
 import pytest
 
 
-def make_ws(env, nodes: dict[str, dict], edges: list[tuple[str, str, str, str]]):
+def make_ws(env, nodes: dict[str, dict], edges: list[tuple[str, str]]):
     ws = workspace.Workspace(env=env)
     for id, data in nodes.items():
+        title = data["title"]
+        del data["title"]
         ws.nodes.append(
             workspace.WorkspaceNode(
                 id=id,
                 type="basic",
-                data=workspace.WorkspaceNodeData(title=data["title"], params=data),
+                data=workspace.WorkspaceNodeData(title=title, params=data),
                 position=workspace.Position(
                     x=data.get("x", 0),
                     y=data.get("y", 0),
@@ -31,35 +33,86 @@ def make_ws(env, nodes: dict[str, dict], edges: list[tuple[str, str, str, str]])
     return ws
 
 
+def summarize_layers(m: pytorch_model_ops.ModelConfig) -> str:
+    return "".join(str(e)[0] for e in m.model)
+
+
+def summarize_connections(m: pytorch_model_ops.ModelConfig) -> str:
+    return " ".join(
+        "".join(n[0] for n in c.param_names) + "->" + "".join(n[0] for n in c.return_names)
+        for c in m.model._children
+    )
+
+
 async def test_build_model():
     ws = make_ws(
         pytorch_model_ops.ENV,
         {
-            "emb": {"title": "Input: embedding"},
-            "lin": {"title": "Linear", "output_dim": "same"},
-            "act": {"title": "Activation", "type": "Leaky ReLU"},
-            "label": {"title": "Input: label"},
+            "emb": {"title": "Input: tensor"},
+            "lin": {"title": "Linear", "output_dim": 4},
+            "act": {"title": "Activation", "type": "Leaky_ReLU"},
+            "label": {"title": "Input: tensor"},
             "loss": {"title": "MSE loss"},
             "optim": {"title": "Optimizer", "type": "SGD", "lr": 0.1},
         },
         [
-            ("emb:x", "lin:x"),
-            ("lin:x", "act:x"),
-            ("act:x", "loss:x"),
-            ("label:y", "loss:y"),
-            ("loss:loss", "optim:loss"),
+            ("emb:output", "lin:x"),
+            ("lin:output", "act:x"),
+            ("act:output", "loss:x"),
+            ("label:output", "loss:y"),
+            ("loss:output", "optim:loss"),
         ],
     )
     x = torch.rand(100, 4)
     y = x + 1
-    m = pytorch_model_ops.build_model(ws, {"emb_x": x, "label_y": y})
+    m = pytorch_model_ops.build_model(ws)
     for i in range(1000):
-        loss = m.train({"emb_x": x, "label_y": y})
+        loss = m.train({"emb_output": x, "label_output": y})
     assert loss < 0.1
-    o = m.inference({"emb_x": x[:1]})
-    error = torch.nn.functional.mse_loss(o["act_x"], x[:1] + 1)
+    o = m.inference({"emb_output": x[:1]})
+    error = torch.nn.functional.mse_loss(o["act_output"], x[:1] + 1)
     assert error < 0.1
 
 
+async def test_build_model_with_repeat():
+    def repeated_ws(times):
+        return make_ws(
+            pytorch_model_ops.ENV,
+            {
+                "emb": {"title": "Input: tensor"},
+                "lin": {"title": "Linear", "output_dim": 8},
+                "act": {"title": "Activation", "type": "Leaky_ReLU"},
+                "label": {"title": "Input: tensor"},
+                "loss": {"title": "MSE loss"},
+                "optim": {"title": "Optimizer", "type": "SGD", "lr": 0.1},
+                "repeat": {"title": "Repeat", "times": times, "same_weights": False},
+            },
+            [
+                ("emb:output", "lin:x"),
+                ("lin:output", "act:x"),
+                ("act:output", "loss:x"),
+                ("label:output", "loss:y"),
+                ("loss:output", "optim:loss"),
+                ("repeat:output", "lin:x"),
+                ("act:output", "repeat:input"),
+            ],
+        )
+
+    # 1 repetition
+    m = pytorch_model_ops.build_model(repeated_ws(1))
+    assert summarize_layers(m) == "IL<II"
+    assert summarize_connections(m) == "e->S S->l l->a a->E E->E"
+
+    # 2 repetitions
+    m = pytorch_model_ops.build_model(repeated_ws(2))
+    assert summarize_layers(m) == "IL<IL<II"
+    assert summarize_connections(m) == "e->S S->l l->a a->S S->l l->a a->E E->E"
+
+    # 3 repetitions
+    m = pytorch_model_ops.build_model(repeated_ws(3))
+    assert summarize_layers(m) == "IL<IL<IL<II"
+    assert summarize_connections(m) == "e->S S->l l->a a->S S->l l->a a->S S->l l->a a->E E->E"
+
+
 if __name__ == "__main__":
     pytest.main()