Merge pull request #119 from biggraph/darabos-more-pytorch

examples/Model definition

@@ -8,45 +8,66 @@
       "targetHandle": "loss"
     },
     {
-      "id": "Activation 1 MSE loss 2",
+      "id": "Activation 1 Repeat 1",
       "source": "Activation 1",
       "sourceHandle": "output",
-      "target": "MSE loss 2",
-      "targetHandle": "x"
+      "target": "Repeat 1",
+      "targetHandle": "input"
     },
     {
-      "id": "Input: tensor 3 MSE loss 2",
-      "source": "Input: tensor 3",
-      "sourceHandle": "x",
-      "target": "MSE loss 2",
-      "targetHandle": "y"
+      "id": "Linear 1 Activation 1",
+      "source": "Linear 1",
+      "sourceHandle": "output",
+      "target": "Activation 1",
+      "targetHandle": "x"
     },
     {
-      "id": "Activation 1 Repeat 1",
-      "source": "Activation 1",
+      "id": "Repeat 1 Linear 1",
+      "source": "Repeat 1",
       "sourceHandle": "output",
-      "target": "Repeat 1",
-      "targetHandle": "input"
+      "target": "Linear 1",
+      "targetHandle": "x"
     },
     {
       "id": "Input: tensor 1 Linear 1",
       "source": "Input: tensor 1",
-      "sourceHandle": "x",
+      "sourceHandle": "output",
       "target": "Linear 1",
       "targetHandle": "x"
     },
     {
-      "id": "Linear 1 Activation 1",
-      "source": "Linear 1",
+      "id": "Constant vector 1 Add 1",
+      "source": "Constant vector 1",
       "sourceHandle": "output",
-      "target": "Activation 1",
-      "targetHandle": "x"
+      "target": "Add 1",
+      "targetHandle": "b"
     },
     {
-      "id": "Repeat 1 Linear 1",
-      "source": "Repeat 1",
+      "id": "Input: tensor 3 Add 1",
+      "source": "Input: tensor 3",
       "sourceHandle": "output",
-      "target": "Linear 1",
+      "target": "Add 1",
+      "targetHandle": "a"
+    },
+    {
+      "id": "Add 1 MSE loss 2",
+      "source": "Add 1",
+      "sourceHandle": "output",
+      "target": "MSE loss 2",
+      "targetHandle": "y"
+    },
+    {
+      "id": "Activation 1 Output 1",
+      "source": "Activation 1",
+      "sourceHandle": "output",
+      "target": "Output 1",
+      "targetHandle": "x"
+    },
+    {
+      "id": "Output 1 MSE loss 2",
+      "source": "Output 1",
+      "sourceHandle": "x",
+      "target": "MSE loss 2",
       "targetHandle": "x"
     }
   ],
@@ -108,8 +129,8 @@
       "height": 250.0,
       "id": "Optimizer 2",
       "position": {
-        "x": 292.3983313429414,
-        "y": -853.8015246037802
+        "x": 359.75221367487865,
+        "y": -1560.7604266065723
       },
       "type": "basic",
       "width": 232.0
@@ -155,10 +176,6 @@
               }
             }
           },
-          "position": {
-            "x": 344.0,
-            "y": 384.0
-          },
           "type": "basic"
         },
         "params": {
@@ -188,58 +205,8 @@
           "inputs": {},
           "name": "Input: tensor",
           "outputs": {
-            "x": {
-              "name": "x",
-              "position": "top",
-              "type": {
-                "type": "tensor"
-              }
-            }
-          },
-          "params": {
-            "name": {
-              "default": null,
-              "name": "name",
-              "type": {
-                "type": "None"
-              }
-            }
-          },
-          "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": {
-            "x": {
-              "name": "x",
+            "output": {
+              "name": "output",
               "position": "top",
               "type": {
                 "type": "tensor"
@@ -255,10 +222,6 @@
               }
             }
           },
-          "position": {
-            "x": 1169.0,
-            "y": 340.0
-          },
           "type": "basic"
         },
         "params": {
@@ -279,6 +242,8 @@
     },
     {
       "data": {
+        "__execution_delay": null,
+        "collapsed": true,
         "display": null,
         "error": null,
         "input_metadata": null,
@@ -310,10 +275,6 @@
             }
           },
           "params": {},
-          "position": {
-            "x": 937.0,
-            "y": 270.0
-          },
           "type": "basic"
         },
         "params": {},
@@ -324,8 +285,8 @@
       "height": 200.0,
       "id": "MSE loss 2",
       "position": {
-        "x": 309.4422414664647,
-        "y": -552.1056805642488
+        "x": 362.77557479979805,
+        "y": -1287.1056805642488
       },
       "type": "basic",
       "width": 200.0
@@ -373,10 +334,6 @@
               }
             }
           },
-          "position": {
-            "x": 487.0,
-            "y": 443.0
-          },
           "type": "basic"
         },
         "params": {
@@ -433,8 +390,8 @@
             }
           },
           "position": {
-            "x": 359.0,
-            "y": 310.0
+            "x": 667.0,
+            "y": 432.0
           },
           "type": "basic"
         },
@@ -448,8 +405,219 @@
       "height": 200.0,
       "id": "Linear 1",
       "position": {
-        "x": 88.83370222907377,
-        "y": 48.642890099180136
+        "x": 98.54861342271252,
+        "y": 14.121603973834155
+      },
+      "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": {
+            "output": {
+              "name": "output",
+              "position": "top",
+              "type": {
+                "type": "tensor"
+              }
+            }
+          },
+          "params": {
+            "name": {
+              "default": null,
+              "name": "name",
+              "type": {
+                "type": "None"
+              }
+            }
+          },
+          "position": {
+            "x": 675.0,
+            "y": 499.0
+          },
+          "type": "basic"
+        },
+        "params": {
+          "name": "X"
+        },
+        "status": "planned",
+        "title": "Input: tensor"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 200.0,
+      "id": "Input: tensor 1",
+      "position": {
+        "x": 108.75735538875443,
+        "y": 331.53404347930933
+      },
+      "type": "basic",
+      "width": 200.0
+    },
+    {
+      "data": {
+        "__execution_delay": 0.0,
+        "collapsed": null,
+        "display": null,
+        "error": null,
+        "input_metadata": null,
+        "meta": {
+          "inputs": {},
+          "name": "Constant vector",
+          "outputs": {
+            "output": {
+              "name": "output",
+              "position": "top",
+              "type": {
+                "type": "None"
+              }
+            }
+          },
+          "params": {
+            "size": {
+              "default": 1.0,
+              "name": "size",
+              "type": {
+                "type": "<class 'int'>"
+              }
+            },
+            "value": {
+              "default": 0.0,
+              "name": "value",
+              "type": {
+                "type": "<class 'int'>"
+              }
+            }
+          },
+          "position": {
+            "x": 1061.0,
+            "y": 239.0
+          },
+          "type": "basic"
+        },
+        "params": {
+          "size": "1",
+          "value": "1"
+        },
+        "status": "planned",
+        "title": "Constant vector"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 258.0,
+      "id": "Constant vector 1",
+      "position": {
+        "x": 983.1241140187901,
+        "y": -562.803650462906
+      },
+      "type": "basic",
+      "width": 238.0
+    },
+    {
+      "data": {
+        "__execution_delay": null,
+        "collapsed": true,
+        "display": null,
+        "error": null,
+        "input_metadata": null,
+        "meta": {
+          "inputs": {
+            "a": {
+              "name": "a",
+              "position": "bottom",
+              "type": {
+                "type": "<class 'inspect._empty'>"
+              }
+            },
+            "b": {
+              "name": "b",
+              "position": "bottom",
+              "type": {
+                "type": "<class 'inspect._empty'>"
+              }
+            }
+          },
+          "name": "Add",
+          "outputs": {
+            "output": {
+              "name": "output",
+              "position": "top",
+              "type": {
+                "type": "None"
+              }
+            }
+          },
+          "params": {},
+          "position": {
+            "x": 1077.0,
+            "y": 334.0
+          },
+          "type": "basic"
+        },
+        "params": {},
+        "status": "planned",
+        "title": "Add"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 200.0,
+      "id": "Add 1",
+      "position": {
+        "x": 818.5444381090571,
+        "y": -955.5157374399466
+      },
+      "type": "basic",
+      "width": 200.0
+    },
+    {
+      "data": {
+        "__execution_delay": null,
+        "collapsed": true,
+        "display": null,
+        "error": null,
+        "input_metadata": null,
+        "meta": {
+          "inputs": {
+            "x": {
+              "name": "x",
+              "position": "bottom",
+              "type": {
+                "type": "tensor"
+              }
+            }
+          },
+          "name": "Output",
+          "outputs": {
+            "x": {
+              "name": "x",
+              "position": "top",
+              "type": {
+                "type": "tensor"
+              }
+            }
+          },
+          "params": {},
+          "position": {
+            "x": 544.0,
+            "y": 306.0
+          },
+          "type": "basic"
+        },
+        "params": {},
+        "status": "planned",
+        "title": "Output"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 200.0,
+      "id": "Output 1",
+      "position": {
+        "x": 185.15239170944702,
+        "y": -733.1526319565451
       },
       "type": "basic",
       "width": 200.0

examples/Model use

@@ -579,54 +579,54 @@
               ],
               "data": [
                 [
-                  "[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]",
+                  "[2.5075035095214844, 3.0133981704711914, 2.698194980621338, 2.3600802421569824]"
                 ],
                 [
-                  "[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.98033333 0.97656083 0.38939917 0.81491041]",
+                  "[1.98033333 1.97656083 1.38939917 1.81491041]",
+                  "[3.0080792903900146, 2.9657773971557617, 2.4899187088012695, 2.7300877571105957]"
                 ],
                 [
-                  "[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.74064726 0.4155122  0.09800029 0.49930882]",
+                  "[1.74064732 1.4155122  1.09800029 1.49930882]",
+                  "[2.7571821212768555, 2.4066381454467773, 2.099902629852295, 2.491953134536743]"
                 ],
                 [
-                  "[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.68891573 0.25576538 0.96339929 0.503833  ]",
+                  "[1.68891573 1.25576544 1.96339929 1.50383306]",
+                  "[2.4299726486206055, 2.4633498191833496, 2.669276237487793, 2.753784656524658]"
                 ],
                 [
-                  "[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.87285906 0.48354989 0.39394957 0.59456545]",
+                  "[1.872859   1.48354983 1.39394951 1.59456539]",
+                  "[2.7922420501708984, 2.555819034576416, 2.3168320655822754, 2.6561851501464844]"
                 ],
                 [
-                  "[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.32565445 0.90939188 0.07488042 0.13730896]",
+                  "[1.32565451 1.90939188 1.07488036 1.13730896]",
+                  "[2.5029656887054443, 2.744194984436035, 2.2190115451812744, 2.024731159210205]"
                 ],
                 [
-                  "[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.49607176 0.1922397  0.46640229 0.78321403]",
+                  "[1.49607182 1.19223976 1.46640229 1.78321409]",
+                  "[2.4844861030578613, 2.2119431495666504, 2.451040267944336, 2.7876930236816406]"
                 ],
                 [
-                  "[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.85566247 0.83362883 0.48424995 0.25265992]",
+                  "[1.85566247 1.83362889 1.48424995 1.25265992]",
+                  "[2.757845401763916, 2.9018521308898926, 2.371169328689575, 2.356513500213623]"
                 ],
                 [
-                  "[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.48959708 0.48549271 0.32688856 0.356677  ]",
+                  "[1.48959708 1.48549271 1.32688856 1.35667706]",
+                  "[2.4994757175445557, 2.469498634338379, 2.3010873794555664, 2.3796985149383545]"
                 ],
                 [
-                  "[0.34084332 0.73018837 0.54168713 0.91440833]",
-                  "[1.34084332 1.73018837 1.54168713 1.91440833]",
-                  "[1.5488454103469849, 1.4963982105255127, 1.422922968864441, 1.622254490852356]"
+                  "[0.80654246 0.08253473 0.74478531 0.71257162]",
+                  "[1.8065424  1.08253479 1.74478531 1.71257162]",
+                  "[2.5831751823425293, 2.2751128673553467, 2.5047459602355957, 2.9079394340515137]"
                 ]
               ]
             },
@@ -688,10 +688,6 @@
                   "[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]"
@@ -700,6 +696,10 @@
                   "[0.04508126 0.76880038 0.80721325 0.62542385]",
                   "[1.04508126 1.76880038 1.80721331 1.62542391]"
                 ],
+                [
+                  "[0.19908059 0.17570406 0.51475513 0.1893943 ]",
+                  "[1.19908059 1.175704   1.51475513 1.18939424]"
+                ],
                 [
                   "[0.40167677 0.25953674 0.9407078  0.76308483]",
                   "[1.40167677 1.25953674 1.9407078  1.76308489]"
@@ -716,10 +716,6 @@
                   "[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]"
@@ -736,6 +732,10 @@
                   "[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]"
@@ -752,6 +752,10 @@
                   "[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]"
@@ -764,14 +768,6 @@
                   "[0.4827103  0.10563457 0.98858833 0.82286644]",
                   "[1.48271036 1.10563457 1.98858833 1.82286644]"
                 ],
-                [
-                  "[0.98033333 0.97656083 0.38939917 0.81491041]",
-                  "[1.98033333 1.97656083 1.38939917 1.81491041]"
-                ],
-                [
-                  "[0.74064726 0.4155122  0.09800029 0.49930882]",
-                  "[1.74064732 1.4155122  1.09800029 1.49930882]"
-                ],
                 [
                   "[0.78956431 0.87284744 0.06880784 0.03455889]",
                   "[1.78956437 1.87284744 1.06880784 1.03455889]"
@@ -808,6 +804,10 @@
                   "[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]"
@@ -836,14 +836,18 @@
                   "[0.18720162 0.74115586 0.98626411 0.30355608]",
                   "[1.18720162 1.74115586 1.98626411 1.30355608]"
                 ],
-                [
-                  "[0.85566247 0.83362883 0.48424995 0.25265992]",
-                  "[1.85566247 1.83362889 1.48424995 1.25265992]"
-                ],
                 [
                   "[0.95928186 0.84273899 0.71514636 0.38619852]",
                   "[1.95928192 1.84273899 1.7151463  1.38619852]"
                 ],
+                [
+                  "[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]"
@@ -852,10 +856,6 @@
                   "[0.67418337 0.79634351 0.23229051 0.71345252]",
                   "[1.67418337 1.79634356 1.23229051 1.71345258]"
                 ],
-                [
-                  "[0.87285906 0.48354989 0.39394957 0.59456545]",
-                  "[1.872859   1.48354983 1.39394951 1.59456539]"
-                ],
                 [
                   "[0.81788456 0.58174163 0.29376316 0.7971254 ]",
                   "[1.81788456 1.58174157 1.29376316 1.79712534]"
@@ -868,6 +868,10 @@
                   "[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]"
@@ -896,10 +900,6 @@
                   "[0.60609657 0.96257663 0.19292736 0.95702219]",
                   "[1.60609651 1.96257663 1.19292736 1.95702219]"
                 ],
-                [
-                  "[0.80654246 0.08253473 0.74478531 0.71257162]",
-                  "[1.8065424  1.08253479 1.74478531 1.71257162]"
-                ],
                 [
                   "[0.70167565 0.26930219 0.5660674  0.61194974]",
                   "[1.70167565 1.26930213 1.56606746 1.61194968]"
@@ -940,14 +940,14 @@
                   "[0.37959969 0.42820001 0.10690689 0.96353984]",
                   "[1.37959969 1.42820001 1.10690689 1.96353984]"
                 ],
-                [
-                  "[0.49607176 0.1922397  0.46640229 0.78321403]",
-                  "[1.49607182 1.19223976 1.46640229 1.78321409]"
-                ],
                 [
                   "[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]"
@@ -980,6 +980,10 @@
                   "[0.68094063 0.45189077 0.22661722 0.37354094]",
                   "[1.68094063 1.45189071 1.22661722 1.37354088]"
                 ],
+                [
+                  "[0.43681622 0.74680805 0.83598751 0.12414402]",
+                  "[1.43681622 1.74680805 1.83598757 1.12414408]"
+                ],
                 [
                   "[0.47870928 0.17129105 0.27300501 0.20634609]",
                   "[1.47870922 1.17129111 1.27300501 1.20634604]"
@@ -991,16 +995,12 @@
                 [
                   "[0.87608397 0.93200487 0.80169648 0.37758952]",
                   "[1.87608397 1.93200493 1.80169654 1.37758946]"
-                ],
-                [
-                  "[0.68891573 0.25576538 0.96339929 0.503833  ]",
-                  "[1.68891573 1.25576544 1.96339929 1.50383306]"
                 ]
               ]
             }
           },
           "other": {
-            "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)"
+            "model": "ModelConfig(model=Sequential(\n  (0) - Identity(): Input__tensor_1_output -> START_Repeat_1_output\n  (1) - Linear(4, 4, bias=True): START_Repeat_1_output -> Linear_1_output\n  (2) - <function leaky_relu at 0x7b4a644509a0>: Linear_1_output -> Activation_1_output\n  (3) - Identity(): Activation_1_output -> START_Repeat_1_output\n  (4) - Linear(4, 4, bias=True): START_Repeat_1_output -> Linear_1_output\n  (5) - <function leaky_relu at 0x7b4a644509a0>: Linear_1_output -> Activation_1_output\n  (6) - Identity(): Activation_1_output -> END_Repeat_1_output\n  (7) - Identity(): END_Repeat_1_output -> Output_1_x\n  (8) - Identity(): Output_1_x -> Output_1_x\n), model_inputs=['Input__tensor_1_output'], model_outputs=['Output_1_x'], loss_inputs=['Input__tensor_3_output', 'Output_1_x'], loss=Sequential(\n  (0) - <function constant_vector.<locals>.<lambda> at 0x7b4975350e00>: nothing -> Constant_vector_1_output\n  (1) - <built-in method add of type object at 0x7b4a5859ef00>: Input__tensor_3_output, Constant_vector_1_output -> Add_1_output\n  (2) - <function mse_loss at 0x7b4a64452480>: Output_1_x, Add_1_output -> MSE_loss_2_output\n  (3) - 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": []
         },
@@ -1032,14 +1032,14 @@
               "model": {
                 "model": {
                   "inputs": [
-                    "Input__tensor_1_x"
+                    "Input__tensor_1_output"
                   ],
                   "loss_inputs": [
-                    "END_Repeat_1_output",
-                    "Input__tensor_3_x"
+                    "Input__tensor_3_output",
+                    "Output_1_x"
                   ],
                   "outputs": [
-                    "END_Repeat_1_output"
+                    "Output_1_x"
                   ],
                   "trained": true
                 },
@@ -1207,14 +1207,14 @@
               "model": {
                 "model": {
                   "inputs": [
-                    "Input__tensor_1_x"
+                    "Input__tensor_1_output"
                   ],
                   "loss_inputs": [
-                    "END_Repeat_1_output",
-                    "Input__tensor_3_x"
+                    "Input__tensor_3_output",
+                    "Output_1_x"
                   ],
                   "outputs": [
-                    "END_Repeat_1_output"
+                    "Output_1_x"
                   ],
                   "trained": false
                 },
@@ -1270,8 +1270,8 @@
           "type": "basic"
         },
         "params": {
-          "epochs": "1003",
-          "input_mapping": "{\"map\":{\"Input__tensor_1_x\":{\"df\":\"df_train\",\"column\":\"x\"},\"Input__tensor_3_x\":{\"df\":\"df_train\",\"column\":\"y\"}}}",
+          "epochs": "1001",
+          "input_mapping": "{\"map\":{\"Input__tensor_1_output\":{\"df\":\"df_train\",\"column\":\"x\"},\"Input__tensor_3_output\":{\"df\":\"df_train\",\"column\":\"y\"}}}",
           "model_name": "model"
         },
         "status": "done",
@@ -1319,14 +1319,14 @@
               "model": {
                 "model": {
                   "inputs": [
-                    "Input__tensor_1_x"
+                    "Input__tensor_1_output"
                   ],
                   "loss_inputs": [
-                    "END_Repeat_1_output",
-                    "Input__tensor_3_x"
+                    "Input__tensor_3_output",
+                    "Output_1_x"
                   ],
                   "outputs": [
-                    "END_Repeat_1_output"
+                    "Output_1_x"
                   ],
                   "trained": true
                 },
@@ -1382,9 +1382,9 @@
           "type": "basic"
         },
         "params": {
-          "input_mapping": "{\"map\":{\"Input__tensor_1_x\":{\"df\":\"df_test\",\"column\":\"x\"}}}",
+          "input_mapping": "{\"map\":{\"Input__tensor_1_output\":{\"df\":\"df_test\",\"column\":\"x\"}}}",
           "model_name": "model",
-          "output_mapping": "{\"map\":{\"END_Repeat_1_output\":{\"df\":\"df_test\",\"column\":\"pred\"}}}"
+          "output_mapping": "{\"map\":{\"Output_1_x\":{\"df\":\"df_test\",\"column\":\"pred\"}}}"
         },
         "status": "done",
         "title": "Model inference"

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

@@ -13,7 +13,7 @@ pd.options.mode.copy_on_write = True  # Prepare for Pandas 3.0.
 from .core import *  # noqa (easier access for core classes)
 from . import lynxkite_ops  # noqa (imported to trigger registration)
 from . import networkx_ops  # noqa (imported to trigger registration)
-from . import pytorch_model_ops  # noqa (imported to trigger registration)
+from . import pytorch  # noqa (imported to trigger registration)
 
 if os.environ.get("LYNXKITE_BIONEMO_INSTALLED", "").strip().lower() == "true":
     from . import bionemo_ops  # noqa (imported to trigger registration)

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

@@ -8,7 +8,8 @@ from lynxkite.core import ops
 from collections import deque
 
 from tqdm import tqdm
-from . import core, pytorch_model_ops
+from . import core
+from .pytorch import pytorch_core
 from lynxkite.core import workspace
 import grandcypher
 import joblib
@@ -347,7 +348,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_core.build_model(ws)
     m.source_workspace = model_workspace
     bundle = bundle.copy()
     bundle.other[save_as] = m
@@ -356,15 +357,15 @@ def define_model(
 
 # These contain the same mapping, but they get different UIs.
 # For inputs, you select existing columns. For outputs, you can create new columns.
-class ModelInferenceInputMapping(pytorch_model_ops.ModelMapping):
+class ModelInferenceInputMapping(pytorch_core.ModelMapping):
     pass
 
 
-class ModelTrainingInputMapping(pytorch_model_ops.ModelMapping):
+class ModelTrainingInputMapping(pytorch_core.ModelMapping):
     pass
 
 
-class ModelOutputMapping(pytorch_model_ops.ModelMapping):
+class ModelOutputMapping(pytorch_core.ModelMapping):
     pass
 
 
@@ -379,7 +380,7 @@ 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)
+    inputs = pytorch_core.to_tensors(bundle, input_mapping)
     t = tqdm(range(epochs), desc="Training model")
     losses = []
     for _ in t:
@@ -406,7 +407,7 @@ def model_inference(
         return ops.Result(bundle, error="Mapping is unset.")
     m = bundle.other[model_name]
     assert m.trained, "The model is not trained."
-    inputs = pytorch_model_ops.to_tensors(bundle, input_mapping)
+    inputs = pytorch_core.to_tensors(bundle, input_mapping)
     outputs = m.inference(inputs)
     bundle = bundle.copy()
     copied = set()

lynxkite-graph-analytics/src/lynxkite_graph_analytics/pytorch/__init__.py

@@ -0,0 +1,2 @@
+from . import pytorch_core  # noqa
+from . import pytorch_ops  # noqa

lynxkite-graph-analytics/src/lynxkite_graph_analytics/{pytorch_model_ops.py → pytorch/pytorch_core.py}

@@ -1,16 +1,14 @@
-"""Boxes for defining PyTorch models."""
+"""Infrastructure for defining PyTorch models."""
 
 import copy
-import enum
 import graphlib
 
 import pydantic
 from lynxkite.core import ops, workspace
-from lynxkite.core.ops import Parameter as P
 import torch
 import torch_geometric.nn as pyg_nn
 import dataclasses
-from . import core
+from .. import core
 
 ENV = "PyTorch model"
 
@@ -42,117 +40,6 @@ def reg(name, inputs=[], outputs=None, params=[]):
     )
 
 
-reg("Input: tensor", outputs=["output"], params=[P.basic("name")])
-reg("Input: graph edges", outputs=["edges"])
-reg("Input: sequential", outputs=["y"])
-
-reg("LSTM", inputs=["x", "h"], outputs=["x", "h"])
-reg(
-    "Neural ODE",
-    inputs=["x"],
-    params=[
-        P.basic("relative_tolerance"),
-        P.basic("absolute_tolerance"),
-        P.options(
-            "method",
-            [
-                "dopri8",
-                "dopri5",
-                "bosh3",
-                "fehlberg2",
-                "adaptive_heun",
-                "euler",
-                "midpoint",
-                "rk4",
-                "explicit_adams",
-                "implicit_adams",
-            ],
-        ),
-    ],
-)
-
-
-reg("Attention", inputs=["q", "k", "v"], outputs=["x", "weights"])
-reg("LayerNorm", inputs=["x"])
-reg("Dropout", inputs=["x"], params=[P.basic("p", 0.5)])
-
-
-@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",
-    inputs=["x", "edges"],
-    outputs=["x"],
-    params=[P.options("type", ["GCNConv", "GATConv", "GATv2Conv", "SAGEConv"])],
-)
-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("Triplet margin loss", inputs=["x", "x_pos", "x_neg"], outputs=["loss"])
-reg("Cross-entropy loss", inputs=["x", "y"], outputs=["loss"])
-reg(
-    "Optimizer",
-    inputs=["loss"],
-    outputs=[],
-    params=[
-        P.options(
-            "type",
-            [
-                "AdamW",
-                "Adafactor",
-                "Adagrad",
-                "SGD",
-                "Lion",
-                "Paged AdamW",
-                "Galore AdamW",
-            ],
-        ),
-        P.basic("lr", 0.001),
-    ],
-)
-
-ops.register_passive_op(
-    ENV,
-    "Repeat",
-    inputs=[ops.Input(name="input", position="top", type="tensor")],
-    outputs=[ops.Output(name="output", position="bottom", type="tensor")],
-    params=[
-        ops.Parameter.basic("times", 1, int),
-        ops.Parameter.basic("same_weights", False, bool),
-    ],
-)
-
-ops.register_passive_op(
-    ENV,
-    "Recurrent chain",
-    inputs=[ops.Input(name="input", position="top", type="tensor")],
-    outputs=[ops.Output(name="output", position="bottom", type="tensor")],
-    params=[],
-)
-
-
 def _to_id(*strings: str) -> str:
     """Replaces all non-alphanumeric characters with underscores."""
     return "_".join("".join(c if c.isalnum() else "_" for c in s) for s in strings)
@@ -168,7 +55,10 @@ class Layer:
     outputs: list[str]
 
     def for_sequential(self):
-        inputs = ", ".join(self.inputs)
+        """The layer signature for pyg.nn.Sequential."""
+        # "nothing" is used as a bogus input if an operation has no inputs.
+        # The module in turn needs to take one argument, but it will always be None.
+        inputs = ", ".join(self.inputs) or "nothing"
         outputs = ", ".join(self.outputs)
         return self.module, f"{inputs} -> {outputs}"
 
@@ -201,8 +91,7 @@ class ModelConfig:
         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)
+        output = self.model(nothing=None, **inputs)
         if not isinstance(output, tuple):
             output = (output,)
         values = {k: v for k, v in zip(self.model_outputs, output)}
@@ -220,8 +109,7 @@ class ModelConfig:
         self.optimizer.zero_grad()
         values = self._forward(inputs)
         values.update(inputs)
-        loss_inputs = [values[i] for i in self.loss_inputs]
-        loss = self.loss(*loss_inputs)
+        loss = self.loss(nothing=None, **values)
         loss.backward()
         self.optimizer.step()
         return loss.item()
@@ -317,20 +205,20 @@ class ModelBuilder:
             # 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")
-            ]
+            [(last, lasth)] = self.in_edges[repeat]["input"]
             del self.dependencies[repeat]
             self.dependencies[end_id] = [last]
-            self.dependencies[real_output].append(end_id)
+            real_edges = [e for e in self.out_edges[last][lasth] if e != (repeat, "input")]
             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.out_edges[end_id] = {"output": []}  # Populated below.
+            for real_output, real_outputh in real_edges:
+                self.dependencies[real_output].append(end_id)
+                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.out_edges[end_id]["output"].append((real_output, real_outputh))
         self.inv_dependencies = {n: [] for n in self.nodes}
         for k, v in self.dependencies.items():
             for i in v:
@@ -429,18 +317,19 @@ class ModelBuilder:
 
     def get_config(self) -> ModelConfig:
         # Split the design into model and loss.
-        loss_nodes = set()
+        model_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)
+            if self.nodes[node_id].data.title == "Output":
+                model_nodes.add(node_id)
+                model_nodes |= self.all_upstream(node_id)
+        assert model_nodes, "The model definition must have at least one Output node."
         layers = []
         loss_layers = []
         for layer in self.layers:
-            if layer.origin_id in loss_nodes:
-                loss_layers.append(layer)
-            else:
+            if layer.origin_id in model_nodes:
                 layers.append(layer)
+            else:
+                loss_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)

lynxkite-graph-analytics/src/lynxkite_graph_analytics/pytorch/pytorch_ops.py

@@ -0,0 +1,183 @@
+"""Boxes for defining PyTorch models."""
+
+import enum
+from lynxkite.core import ops
+from lynxkite.core.ops import Parameter as P
+import torch
+import torch_geometric.nn as pyg_nn
+from .pytorch_core import op, reg, ENV
+
+reg("Input: tensor", outputs=["output"], params=[P.basic("name")])
+reg("Input: graph edges", outputs=["edges"])
+reg("Input: sequential", outputs=["y"])
+reg("Output", inputs=["x"], outputs=["x"], params=[P.basic("name")])
+
+
+@op("LSTM")
+def lstm(x, *, input_size=1024, hidden_size=1024, dropout=0.0):
+    return torch.nn.LSTM(input_size, hidden_size, dropout=0.0)
+
+
+reg(
+    "Neural ODE",
+    inputs=["x"],
+    params=[
+        P.basic("relative_tolerance"),
+        P.basic("absolute_tolerance"),
+        P.options(
+            "method",
+            [
+                "dopri8",
+                "dopri5",
+                "bosh3",
+                "fehlberg2",
+                "adaptive_heun",
+                "euler",
+                "midpoint",
+                "rk4",
+                "explicit_adams",
+                "implicit_adams",
+            ],
+        ),
+    ],
+)
+
+
+@op("Attention", outputs=["outputs", "weights"])
+def attention(query, key, value, *, embed_dim=1024, num_heads=1, dropout=0.0):
+    return torch.nn.MultiHeadAttention(embed_dim, num_heads, dropout=dropout, need_weights=True)
+
+
+@op("LayerNorm", outputs=["outputs", "weights"])
+def layernorm(x, *, normalized_shape=""):
+    normalized_shape = [int(s.strip()) for s in normalized_shape.split(",")]
+    return torch.nn.LayerNorm(normalized_shape)
+
+
+@op("Dropout", outputs=["outputs", "weights"])
+def dropout(x, *, p=0.0):
+    return torch.nn.Dropout(p)
+
+
+@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
+
+
+@op("Constant vector")
+def constant_vector(*, value=0, size=1):
+    return lambda _: torch.full((size,), value)
+
+
+@op("Softmax")
+def softmax(x, *, dim=1):
+    return torch.nn.Softmax(dim=dim)
+
+
+@op("Concatenate")
+def concatenate(a, b):
+    return lambda a, b: torch.concatenate(*torch.broadcast_tensors(a, b))
+
+
+reg(
+    "Graph conv",
+    inputs=["x", "edges"],
+    outputs=["x"],
+    params=[P.options("type", ["GCNConv", "GATConv", "GATv2Conv", "SAGEConv"])],
+)
+
+reg("Triplet margin loss", inputs=["x", "x_pos", "x_neg"], outputs=["loss"])
+reg("Cross-entropy loss", inputs=["x", "y"], outputs=["loss"])
+reg(
+    "Optimizer",
+    inputs=["loss"],
+    outputs=[],
+    params=[
+        P.options(
+            "type",
+            [
+                "AdamW",
+                "Adafactor",
+                "Adagrad",
+                "SGD",
+                "Lion",
+                "Paged AdamW",
+                "Galore AdamW",
+            ],
+        ),
+        P.basic("lr", 0.001),
+    ],
+)
+
+ops.register_passive_op(
+    ENV,
+    "Repeat",
+    inputs=[ops.Input(name="input", position="top", type="tensor")],
+    outputs=[ops.Output(name="output", position="bottom", type="tensor")],
+    params=[
+        ops.Parameter.basic("times", 1, int),
+        ops.Parameter.basic("same_weights", False, bool),
+    ],
+)
+
+ops.register_passive_op(
+    ENV,
+    "Recurrent chain",
+    inputs=[ops.Input(name="input", position="top", type="tensor")],
+    outputs=[ops.Output(name="output", position="bottom", type="tensor")],
+    params=[],
+)
+
+
+def _set_handle_positions(op):
+    op: ops.Op = op.__op__
+    for v in op.outputs.values():
+        v.position = "top"
+    for v in op.inputs.values():
+        v.position = "bottom"
+
+
+def _register_simple_pytorch_layer(func):
+    op = ops.op(ENV, func.__name__.title())(lambda input: func)
+    _set_handle_positions(op)
+
+
+def _register_two_tensor_function(func):
+    op = ops.op(ENV, func.__name__.title())(lambda a, b: func)
+    _set_handle_positions(op)
+
+
+SIMPLE_FUNCTIONS = [
+    torch.sin,
+    torch.cos,
+    torch.log,
+    torch.exp,
+]
+TWO_TENSOR_FUNCTIONS = [
+    torch.multiply,
+    torch.add,
+    torch.subtract,
+]
+
+
+for f in SIMPLE_FUNCTIONS:
+    _register_simple_pytorch_layer(f)
+for f in TWO_TENSOR_FUNCTIONS:
+    _register_two_tensor_function(f)

lynxkite-graph-analytics/tests/test_pytorch_model_ops.py

@@ -1,5 +1,5 @@
 from lynxkite.core import workspace
-from lynxkite_graph_analytics import pytorch_model_ops
+from lynxkite_graph_analytics.pytorch import pytorch_core
 import torch
 import pytest
 
@@ -33,11 +33,11 @@ def make_ws(env, nodes: dict[str, dict], edges: list[tuple[str, str]]):
     return ws
 
 
-def summarize_layers(m: pytorch_model_ops.ModelConfig) -> str:
+def summarize_layers(m: pytorch_core.ModelConfig) -> str:
     return "".join(str(e)[0] for e in m.model)
 
 
-def summarize_connections(m: pytorch_model_ops.ModelConfig) -> str:
+def summarize_connections(m: pytorch_core.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
@@ -46,51 +46,55 @@ def summarize_connections(m: pytorch_model_ops.ModelConfig) -> str:
 
 async def test_build_model():
     ws = make_ws(
-        pytorch_model_ops.ENV,
+        pytorch_core.ENV,
         {
-            "emb": {"title": "Input: tensor"},
+            "input": {"title": "Input: tensor"},
             "lin": {"title": "Linear", "output_dim": 4},
             "act": {"title": "Activation", "type": "Leaky_ReLU"},
+            "output": {"title": "Output"},
             "label": {"title": "Input: tensor"},
             "loss": {"title": "MSE loss"},
             "optim": {"title": "Optimizer", "type": "SGD", "lr": 0.1},
         },
         [
-            ("emb:output", "lin:x"),
+            ("input:output", "lin:x"),
             ("lin:output", "act:x"),
-            ("act:output", "loss:x"),
+            ("act:output", "output:x"),
+            ("output:x", "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)
+    m = pytorch_core.build_model(ws)
     for i in range(1000):
-        loss = m.train({"emb_output": x, "label_output": y})
+        loss = m.train({"input_output": x, "label_output": y})
     assert loss < 0.1
-    o = m.inference({"emb_output": x[:1]})
-    error = torch.nn.functional.mse_loss(o["act_output"], x[:1] + 1)
+    o = m.inference({"input_output": x[:1]})
+    error = torch.nn.functional.mse_loss(o["output_x"], 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,
+            pytorch_core.ENV,
             {
-                "emb": {"title": "Input: tensor"},
+                "input": {"title": "Input: tensor"},
                 "lin": {"title": "Linear", "output_dim": 8},
                 "act": {"title": "Activation", "type": "Leaky_ReLU"},
+                "output": {"title": "Output"},
                 "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"),
+                ("input:output", "lin:x"),
                 ("lin:output", "act:x"),
-                ("act:output", "loss:x"),
+                ("act:output", "output:x"),
+                ("output:x", "loss:x"),
                 ("label:output", "loss:y"),
                 ("loss:output", "optim:loss"),
                 ("repeat:output", "lin:x"),
@@ -99,19 +103,19 @@ async def test_build_model_with_repeat():
         )
 
     # 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"
+    m = pytorch_core.build_model(repeated_ws(1))
+    assert summarize_layers(m) == "IL<III"
+    assert summarize_connections(m) == "i->S S->l l->a a->E E->o o->o"
 
     # 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"
+    m = pytorch_core.build_model(repeated_ws(2))
+    assert summarize_layers(m) == "IL<IL<III"
+    assert summarize_connections(m) == "i->S S->l l->a a->S S->l l->a a->E E->o o->o"
 
     # 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"
+    m = pytorch_core.build_model(repeated_ws(3))
+    assert summarize_layers(m) == "IL<IL<IL<III"
+    assert summarize_connections(m) == "i->S S->l l->a a->S S->l l->a a->S S->l l->a a->E E->o o->o"
 
 
 if __name__ == "__main__":