Each op makes its own layer. #107

examples/Model definition

@@ -1,52 +1,52 @@
 {
   "edges": [
     {
-      "id": "Input: embedding 1 Linear 1",
-      "source": "Input: embedding 1",
-      "sourceHandle": "x",
-      "target": "Linear 1",
-      "targetHandle": "x"
-    },
-    {
-      "id": "Input: label 1 MSE loss 1",
-      "source": "Input: label 1",
-      "sourceHandle": "y",
-      "target": "MSE loss 1",
-      "targetHandle": "y"
+      "id": "MSE loss 1 Optimizer 2",
+      "source": "MSE loss 1",
+      "sourceHandle": "loss",
+      "target": "Optimizer 2",
+      "targetHandle": "loss"
     },
     {
-      "id": "Linear 1 Activation 2",
-      "source": "Linear 1",
-      "sourceHandle": "x",
-      "target": "Activation 2",
+      "id": "Repeat 1 Linear 2",
+      "source": "Repeat 1",
+      "sourceHandle": "output",
+      "target": "Linear 2",
       "targetHandle": "x"
     },
     {
-      "id": "Activation 2 MSE loss 1",
-      "source": "Activation 2",
-      "sourceHandle": "x",
+      "id": "Activation 1 MSE loss 1",
+      "source": "Activation 1",
+      "sourceHandle": "output",
       "target": "MSE loss 1",
       "targetHandle": "x"
     },
     {
-      "id": "MSE loss 1 Optimizer 2",
-      "source": "MSE loss 1",
-      "sourceHandle": "loss",
-      "target": "Optimizer 2",
-      "targetHandle": "loss"
+      "id": "Linear 2 Activation 1",
+      "source": "Linear 2",
+      "sourceHandle": "output",
+      "target": "Activation 1",
+      "targetHandle": "x"
     },
     {
-      "id": "Activation 2 Repeat 1",
-      "source": "Activation 2",
-      "sourceHandle": "x",
+      "id": "Activation 1 Repeat 1",
+      "source": "Activation 1",
+      "sourceHandle": "output",
       "target": "Repeat 1",
       "targetHandle": "input"
     },
     {
-      "id": "Repeat 1 Linear 1",
-      "source": "Repeat 1",
-      "sourceHandle": "output",
-      "target": "Linear 1",
+      "id": "Input: tensor 3 MSE loss 1",
+      "source": "Input: tensor 3",
+      "sourceHandle": "x",
+      "target": "MSE loss 1",
+      "targetHandle": "y"
+    },
+    {
+      "id": "Input: tensor 1 Linear 2",
+      "source": "Input: tensor 1",
+      "sourceHandle": "x",
+      "target": "Linear 2",
       "targetHandle": "x"
     }
   ],
@@ -58,11 +58,26 @@
         "error": null,
         "input_metadata": null,
         "meta": {
-          "inputs": {},
-          "name": "Input: embedding",
-          "outputs": {
+          "inputs": {
             "x": {
               "name": "x",
+              "position": "bottom",
+              "type": {
+                "type": "tensor"
+              }
+            },
+            "y": {
+              "name": "y",
+              "position": "bottom",
+              "type": {
+                "type": "tensor"
+              }
+            }
+          },
+          "name": "MSE loss",
+          "outputs": {
+            "loss": {
+              "name": "loss",
               "position": "top",
               "type": {
                 "type": "tensor"
@@ -74,115 +89,138 @@
         },
         "params": {},
         "status": "planned",
-        "title": "Input: embedding"
+        "title": "MSE loss"
       },
       "dragHandle": ".bg-primary",
       "height": 200.0,
-      "id": "Input: embedding 1",
+      "id": "MSE loss 1",
       "position": {
-        "x": 91.0,
-        "y": 266.0
+        "x": 315.0,
+        "y": -510.0
       },
       "type": "basic",
       "width": 200.0
     },
     {
       "data": {
+        "__execution_delay": 0.0,
+        "collapsed": null,
         "display": null,
         "error": null,
         "input_metadata": null,
         "meta": {
           "inputs": {
-            "x": {
-              "name": "x",
+            "loss": {
+              "name": "loss",
               "position": "bottom",
               "type": {
                 "type": "tensor"
               }
             }
           },
-          "name": "Linear",
-          "outputs": {
-            "x": {
-              "name": "x",
-              "position": "top",
+          "name": "Optimizer",
+          "outputs": {},
+          "params": {
+            "lr": {
+              "default": 0.001,
+              "name": "lr",
               "type": {
-                "type": "tensor"
+                "type": "<class 'float'>"
               }
-            }
-          },
-          "params": {
-            "output_dim": {
-              "default": "same",
-              "name": "output_dim",
+            },
+            "type": {
+              "default": "AdamW",
+              "name": "type",
               "type": {
-                "type": "<class 'str'>"
+                "enum": [
+                  "AdamW",
+                  "Adafactor",
+                  "Adagrad",
+                  "SGD",
+                  "Lion",
+                  "Paged AdamW",
+                  "Galore AdamW"
+                ]
               }
             }
           },
           "type": "basic"
         },
         "params": {
-          "output_dim": "same"
+          "lr": "0.1",
+          "type": "SGD"
         },
         "status": "planned",
-        "title": "Linear"
+        "title": "Optimizer"
       },
       "dragHandle": ".bg-primary",
-      "height": 200.0,
-      "id": "Linear 1",
+      "height": 250.0,
+      "id": "Optimizer 2",
       "position": {
-        "x": 86.0,
-        "y": 33.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",
+            "input": {
+              "name": "input",
+              "position": "top",
               "type": {
                 "type": "tensor"
               }
-            },
-            "y": {
-              "name": "y",
+            }
+          },
+          "name": "Repeat",
+          "outputs": {
+            "output": {
+              "name": "output",
               "position": "bottom",
               "type": {
                 "type": "tensor"
               }
             }
           },
-          "name": "MSE loss",
-          "outputs": {
-            "loss": {
-              "name": "loss",
-              "position": "top",
+          "params": {
+            "same_weights": {
+              "default": false,
+              "name": "same_weights",
               "type": {
-                "type": "tensor"
+                "type": "<class 'bool'>"
+              }
+            },
+            "times": {
+              "default": 1.0,
+              "name": "times",
+              "type": {
+                "type": "<class 'int'>"
               }
             }
           },
-          "params": {},
           "type": "basic"
         },
-        "params": {},
+        "params": {
+          "same_weights": false,
+          "times": "3"
+        },
         "status": "planned",
-        "title": "MSE loss"
+        "title": "Repeat"
       },
       "dragHandle": ".bg-primary",
       "height": 200.0,
-      "id": "MSE loss 1",
+      "id": "Repeat 1",
       "position": {
-        "x": 315.0,
-        "y": -510.0
+        "x": -180.0,
+        "y": -90.0
       },
       "type": "basic",
       "width": 200.0
@@ -193,30 +231,48 @@
         "error": null,
         "input_metadata": null,
         "meta": {
-          "inputs": {},
-          "name": "Input: label",
+          "inputs": {
+            "x": {
+              "name": "x",
+              "position": "bottom",
+              "type": {
+                "type": "<class 'inspect._empty'>"
+              }
+            }
+          },
+          "name": "Linear",
           "outputs": {
-            "y": {
-              "name": "y",
+            "output": {
+              "name": "output",
               "position": "top",
               "type": {
-                "type": "tensor"
+                "type": "None"
+              }
+            }
+          },
+          "params": {
+            "output_dim": {
+              "default": "same",
+              "name": "output_dim",
+              "type": {
+                "type": "<class 'str'>"
               }
             }
           },
-          "params": {},
           "type": "basic"
         },
-        "params": {},
+        "params": {
+          "output_dim": "same"
+        },
         "status": "planned",
-        "title": "Input: label"
+        "title": "Linear"
       },
       "dragHandle": ".bg-primary",
       "height": 200.0,
-      "id": "Input: label 1",
+      "id": "Linear 2",
       "position": {
-        "x": 615.0,
-        "y": -165.0
+        "x": 92.32755761444682,
+        "y": 20.626371289630676
       },
       "type": "basic",
       "width": 200.0
@@ -234,17 +290,17 @@
               "name": "x",
               "position": "bottom",
               "type": {
-                "type": "tensor"
+                "type": "<class 'inspect._empty'>"
               }
             }
           },
           "name": "Activation",
           "outputs": {
-            "x": {
-              "name": "x",
+            "output": {
+              "name": "output",
               "position": "top",
               "type": {
-                "type": "tensor"
+                "type": "None"
               }
             }
           },
@@ -255,27 +311,31 @@
               "type": {
                 "enum": [
                   "ReLU",
-                  "Leaky ReLU",
+                  "Leaky_ReLU",
                   "Tanh",
                   "Mish"
                 ]
               }
             }
           },
+          "position": {
+            "x": 344.0,
+            "y": 384.0
+          },
           "type": "basic"
         },
         "params": {
-          "type": "Leaky ReLU"
+          "type": "Leaky_ReLU"
         },
         "status": "planned",
         "title": "Activation"
       },
       "dragHandle": ".bg-primary",
       "height": 200.0,
-      "id": "Activation 2",
+      "id": "Activation 1",
       "position": {
-        "x": 93.61643829835265,
-        "y": -229.04087132886406
+        "x": 99.77615018185415,
+        "y": -249.43925929074078
       },
       "type": "basic",
       "width": 200.0
@@ -288,56 +348,44 @@
         "error": null,
         "input_metadata": null,
         "meta": {
-          "inputs": {
-            "loss": {
-              "name": "loss",
-              "position": "bottom",
+          "inputs": {},
+          "name": "Input: tensor",
+          "outputs": {
+            "x": {
+              "name": "x",
+              "position": "top",
               "type": {
                 "type": "tensor"
               }
             }
           },
-          "name": "Optimizer",
-          "outputs": {},
           "params": {
-            "lr": {
-              "default": 0.001,
-              "name": "lr",
-              "type": {
-                "type": "<class 'float'>"
-              }
-            },
-            "type": {
-              "default": "AdamW",
-              "name": "type",
+            "name": {
+              "default": null,
+              "name": "name",
               "type": {
-                "enum": [
-                  "AdamW",
-                  "Adafactor",
-                  "Adagrad",
-                  "SGD",
-                  "Lion",
-                  "Paged AdamW",
-                  "Galore AdamW"
-                ]
+                "type": "None"
               }
             }
           },
+          "position": {
+            "x": 258.0,
+            "y": 397.0
+          },
           "type": "basic"
         },
         "params": {
-          "lr": "0.1",
-          "type": "SGD"
+          "name": "X"
         },
         "status": "planned",
-        "title": "Optimizer"
+        "title": "Input: tensor"
       },
       "dragHandle": ".bg-primary",
       "height": 200.0,
-      "id": "Optimizer 2",
+      "id": "Input: tensor 1",
       "position": {
-        "x": 305.6132943499785,
-        "y": -804.0094318451224
+        "x": 97.60681762952905,
+        "y": 293.6278596776366
       },
       "type": "basic",
       "width": 200.0
@@ -350,60 +398,44 @@
         "error": null,
         "input_metadata": null,
         "meta": {
-          "inputs": {
-            "input": {
-              "name": "input",
-              "position": "top",
-              "type": {
-                "type": "tensor"
-              }
-            }
-          },
-          "name": "Repeat",
+          "inputs": {},
+          "name": "Input: tensor",
           "outputs": {
-            "output": {
-              "name": "output",
-              "position": "bottom",
+            "x": {
+              "name": "x",
+              "position": "top",
               "type": {
                 "type": "tensor"
               }
             }
           },
           "params": {
-            "same_weights": {
-              "default": true,
-              "name": "same_weights",
-              "type": {
-                "type": "<class 'bool'>"
-              }
-            },
-            "times": {
-              "default": 1.0,
-              "name": "times",
+            "name": {
+              "default": null,
+              "name": "name",
               "type": {
-                "type": "<class 'int'>"
+                "type": "None"
               }
             }
           },
           "position": {
-            "x": 386.0,
-            "y": 456.0
+            "x": 1169.0,
+            "y": 340.0
           },
           "type": "basic"
         },
         "params": {
-          "same_weights": false,
-          "times": "2"
+          "name": "Y"
         },
         "status": "planned",
-        "title": "Repeat"
+        "title": "Input: tensor"
       },
       "dragHandle": ".bg-primary",
       "height": 200.0,
-      "id": "Repeat 1",
+      "id": "Input: tensor 3",
       "position": {
-        "x": -180.0,
-        "y": -90.0
+        "x": 862.4359094222825,
+        "y": -290.0677203273021
       },
       "type": "basic",
       "width": 200.0

examples/Model use

@@ -579,54 +579,54 @@
               ],
               "data": [
                 [
-                  "[0.67418337 0.79634351 0.23229051 0.71345252]",
-                  "[1.67418337 1.79634356 1.23229051 1.71345258]",
-                  "[1.2304607629776, 1.2535771131515503, 1.4764351844787598, 1.524468183517456]"
+                  "[0.31518555 0.49643308 0.11509258 0.95458382]",
+                  "[1.31518555 1.49643302 1.11509252 1.95458388]",
+                  "[1.3819222450256348, -0.005686390213668346, 1.3793643712997437, 1.581865906715393]"
                 ],
                 [
-                  "[0.6032477  0.83361369 0.18538666 0.19108021]",
-                  "[1.60324764 1.83361363 1.18538666 1.19108021]",
-                  "[1.0745662450790405, 1.3498694896697998, 1.524811029434204, 1.2336101531982422]"
+                  "[0.02162331 0.81861657 0.92468154 0.07808572]",
+                  "[1.02162337 1.81861663 1.92468154 1.07808566]",
+                  "[1.312654972076416, -0.00689137727022171, 1.4941580295562744, 1.243792176246643]"
                 ],
                 [
-                  "[0.62569475 0.9881897  0.83639616 0.9828859 ]",
-                  "[1.62569475 1.9881897  1.83639622 1.98288584]",
-                  "[1.3193368911743164, 1.691685438156128, 1.4516379833221436, 1.6486209630966187]"
+                  "[0.94221359 0.57740951 0.98649532 0.40934443]",
+                  "[1.94221354 1.57740951 1.98649526 1.40934443]",
+                  "[1.9255921840667725, -0.008701151236891747, 1.751355767250061, 1.79597806930542]"
                 ],
                 [
-                  "[0.63235509 0.70352674 0.96188956 0.46240485]",
-                  "[1.63235509 1.70352674 1.96188951 1.46240485]",
-                  "[1.298614740371704, 1.7463937997817993, 1.4211854934692383, 1.2179659605026245]"
+                  "[0.34084332 0.73018837 0.54168713 0.91440833]",
+                  "[1.34084332 1.73018837 1.54168713 1.91440833]",
+                  "[1.6509568691253662, -0.007272087037563324, 1.5942981243133545, 1.81572687625885]"
                 ],
                 [
-                  "[0.74064726 0.4155122  0.09800029 0.49930882]",
-                  "[1.74064732 1.4155122  1.09800029 1.49930882]",
-                  "[1.2740169763565063, 1.0242725610733032, 1.413628339767456, 1.288135290145874]"
+                  "[0.85566247 0.83362883 0.48424995 0.25265992]",
+                  "[1.85566247 1.83362889 1.48424995 1.25265992]",
+                  "[1.7482354640960693, -0.0063837491907179356, 1.4504402875900269, 1.5329445600509644]"
                 ],
                 [
-                  "[0.34084332 0.73018837 0.54168713 0.91440833]",
-                  "[1.34084332 1.73018837 1.54168713 1.91440833]",
-                  "[1.0640922784805298, 1.3728828430175781, 1.1964272260665894, 1.5791122913360596]"
+                  "[0.02235305 0.52774918 0.7331115  0.84358269]",
+                  "[1.02235305 1.52774918 1.7331115  1.84358263]",
+                  "[1.3979142904281616, -0.007555779069662094, 1.6136289834976196, 1.6417407989501953]"
                 ],
                 [
-                  "[0.80654246 0.08253473 0.74478531 0.71257162]",
-                  "[1.8065424  1.08253479 1.74478531 1.71257162]",
-                  "[1.521227240562439, 1.247929334640503, 1.2715831995010376, 1.1901893615722656]"
+                  "[0.9829582  0.59269661 0.40120947 0.95487177]",
+                  "[1.9829582  1.59269667 1.40120947 1.95487177]",
+                  "[1.9523842334747314, -0.00748100271448493, 1.6264307498931885, 1.9942888021469116]"
                 ],
                 [
-                  "[0.11560339 0.57495481 0.76535827 0.0391947 ]",
-                  "[1.11560345 1.57495475 1.76535821 1.0391947 ]",
-                  "[0.8016152381896973, 1.6247310638427734, 1.1153241395950317, 0.9691400527954102]"
+                  "[0.49584109 0.80599248 0.07096875 0.75872749]",
+                  "[1.49584103 1.80599248 1.07096875 1.75872755]",
+                  "[1.5513110160827637, -0.005337317008525133, 1.3384482860565186, 1.5973539352416992]"
                 ],
                 [
-                  "[0.68094063 0.45189077 0.22661722 0.37354094]",
-                  "[1.68094063 1.45189071 1.22661722 1.37354088]",
-                  "[1.224153995513916, 1.1527836322784424, 1.4024348258972168, 1.204542636871338]"
+                  "[0.00497234 0.39319336 0.57054168 0.75150961]",
+                  "[1.00497234 1.39319336 1.57054162 1.75150967]",
+                  "[1.2277441024780273, -0.0067505668848752975, 1.4969637393951416, 1.4524610042572021]"
                 ],
                 [
-                  "[0.79121011 0.54161114 0.69369799 0.1520769 ]",
-                  "[1.79121017 1.54161119 1.69369793 1.15207696]",
-                  "[1.3464093208312988, 1.5594894886016846, 1.5191831588745117, 1.0183898210525513]"
+                  "[0.59492421 0.90274489 0.38069052 0.46101224]",
+                  "[1.59492421 1.90274489 1.38069057 1.46101224]",
+                  "[1.6593225002288818, -0.006088308058679104, 1.4240546226501465, 1.570335865020752]"
                 ]
               ]
             },
@@ -644,6 +644,10 @@
                   "[0.85706753 0.61447072 0.41741937 0.85147089]",
                   "[1.85706758 1.61447072 1.41741943 1.85147095]"
                 ],
+                [
+                  "[0.11560339 0.57495481 0.76535827 0.0391947 ]",
+                  "[1.11560345 1.57495475 1.76535821 1.0391947 ]"
+                ],
                 [
                   "[0.19409031 0.68692201 0.60667384 0.57829887]",
                   "[1.19409037 1.68692207 1.60667384 1.57829881]"
@@ -712,6 +716,10 @@
                   "[0.24388778 0.07268471 0.68350857 0.73431659]",
                   "[1.24388778 1.07268476 1.68350863 1.73431659]"
                 ],
+                [
+                  "[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]"
@@ -733,8 +741,8 @@
                   "[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]"
                 ],
                 [
                   "[0.79423058 0.07138705 0.061777   0.18766576]",
@@ -769,16 +777,12 @@
                   "[1.98033333 1.97656083 1.38939917 1.81491041]"
                 ],
                 [
-                  "[0.78956431 0.87284744 0.06880784 0.03455889]",
-                  "[1.78956437 1.87284744 1.06880784 1.03455889]"
-                ],
-                [
-                  "[0.94221359 0.57740951 0.98649532 0.40934443]",
-                  "[1.94221354 1.57740951 1.98649526 1.40934443]"
+                  "[0.74064726 0.4155122  0.09800029 0.49930882]",
+                  "[1.74064732 1.4155122  1.09800029 1.49930882]"
                 ],
                 [
-                  "[0.00497234 0.39319336 0.57054168 0.75150961]",
-                  "[1.00497234 1.39319336 1.57054162 1.75150967]"
+                  "[0.78956431 0.87284744 0.06880784 0.03455889]",
+                  "[1.78956437 1.87284744 1.06880784 1.03455889]"
                 ],
                 [
                   "[0.44330525 0.09997386 0.89025736 0.90507984]",
@@ -832,10 +836,6 @@
                   "[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]"
@@ -844,10 +844,6 @@
                   "[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 ]"
@@ -856,6 +852,10 @@
                   "[0.54914117 0.03810108 0.87531954 0.73044223]",
                   "[1.54914117 1.03810108 1.87531948 1.73044229]"
                 ],
+                [
+                  "[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]"
@@ -876,14 +876,6 @@
                   "[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]"
-                ],
-                [
-                  "[0.02235305 0.52774918 0.7331115  0.84358269]",
-                  "[1.02235305 1.52774918 1.7331115  1.84358263]"
-                ],
                 [
                   "[0.6080932  0.56563014 0.32107437 0.72599429]",
                   "[1.60809326 1.5656302  1.32107437 1.72599435]"
@@ -904,6 +896,10 @@
                   "[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]"
@@ -912,10 +908,6 @@
                   "[0.76933283 0.86241865 0.44114518 0.65644735]",
                   "[1.76933289 1.86241865 1.44114518 1.65644741]"
                 ],
-                [
-                  "[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]"
@@ -932,6 +924,14 @@
                   "[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]"
+                ],
                 [
                   "[0.37959969 0.42820001 0.10690689 0.96353984]",
                   "[1.37959969 1.42820001 1.10690689 1.96353984]"
@@ -960,10 +960,6 @@
                   "[0.47856545 0.46267092 0.6376707  0.84747767]",
                   "[1.47856545 1.46267092 1.63767076 1.84747767]"
                 ],
-                [
-                  "[0.49584109 0.80599248 0.07096875 0.75872749]",
-                  "[1.49584103 1.80599248 1.07096875 1.75872755]"
-                ],
                 [
                   "[0.43500566 0.66041756 0.80293626 0.96224713]",
                   "[1.43500566 1.66041756 1.80293632 1.96224713]"
@@ -976,6 +972,10 @@
                   "[0.28942841 0.05601001 0.33039129 0.27781558]",
                   "[1.28942847 1.05601001 1.33039129 1.27781558]"
                 ],
+                [
+                  "[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]"
@@ -1000,7 +1000,7 @@
             }
           },
           "other": {
-            "model": "ModelConfig(model=Sequential(\n  (0) - Identity(): Input__embedding_1_x -> START_Repeat_1_output\n  (1) - Linear(in_features=4, out_features=4, bias=True): START_Repeat_1_output -> Linear_1_x\n  (2) - <function leaky_relu at 0x7e6e1cf2c220>: Linear_1_x -> Activation_2_x\n  (3) - Identity(): Activation_2_x -> END_Repeat_1_output\n  (4) - Identity(): END_Repeat_1_output -> END_Repeat_1_output\n), model_inputs=['Input__embedding_1_x'], model_outputs=['END_Repeat_1_output'], loss_inputs=['Input__label_1_y', 'END_Repeat_1_output'], loss=Sequential(\n  (0) - <function mse_loss at 0x7e6e1cf2dd00>: END_Repeat_1_output, 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_2_output\n  (2) - <function leaky_relu at 0x759513340220>: Linear_2_output -> Activation_1_output\n  (3) - Identity(): Activation_1_output -> END_Repeat_1_output\n  (4) - 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 0x759513341d00>: END_Repeat_1_output, Input__tensor_3_x -> 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)"
           },
           "relations": []
         },
@@ -1032,11 +1032,11 @@
               "model": {
                 "model": {
                   "inputs": [
-                    "Input__embedding_1_x"
+                    "Input__tensor_1_x"
                   ],
                   "loss_inputs": [
-                    "Input__label_1_y",
-                    "END_Repeat_1_output"
+                    "END_Repeat_1_output",
+                    "Input__tensor_3_x"
                   ],
                   "outputs": [
                     "END_Repeat_1_output"
@@ -1207,11 +1207,11 @@
               "model": {
                 "model": {
                   "inputs": [
-                    "Input__embedding_1_x"
+                    "Input__tensor_1_x"
                   ],
                   "loss_inputs": [
-                    "Input__label_1_y",
-                    "END_Repeat_1_output"
+                    "END_Repeat_1_output",
+                    "Input__tensor_3_x"
                   ],
                   "outputs": [
                     "END_Repeat_1_output"
@@ -1270,8 +1270,8 @@
           "type": "basic"
         },
         "params": {
-          "epochs": "15",
-          "input_mapping": "{\"map\":{\"Input__embedding_1_x\":{\"column\":\"x\",\"df\":\"df_train\"},\"Input__label_1_y\":{\"column\":\"y\",\"df\":\"df_train\"}}}",
+          "epochs": "150",
+          "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",
@@ -1319,11 +1319,11 @@
               "model": {
                 "model": {
                   "inputs": [
-                    "Input__embedding_1_x"
+                    "Input__tensor_1_x"
                   ],
                   "loss_inputs": [
-                    "Input__label_1_y",
-                    "END_Repeat_1_output"
+                    "END_Repeat_1_output",
+                    "Input__tensor_3_x"
                   ],
                   "outputs": [
                     "END_Repeat_1_output"
@@ -1382,7 +1382,7 @@
           "type": "basic"
         },
         "params": {
-          "input_mapping": "{\"map\":{\"Input__embedding_1_x\":{\"df\":\"df_test\",\"column\":\"x\"}}}",
+          "input_mapping": "{\"map\":{\"Input__tensor_1_x\":{\"df\":\"df_test\",\"column\":\"x\"}}}",
           "model_name": "model",
           "output_mapping": "{\"map\":{\"END_Repeat_1_output\":{\"df\":\"df_test\",\"column\":\"predicted\"}}}"
         },

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

@@ -1,6 +1,7 @@
 """Boxes for defining PyTorch models."""
 
 import copy
+import enum
 import graphlib
 import types
 
@@ -15,6 +16,21 @@ 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 +43,9 @@ def reg(name, inputs=[], outputs=None, params=[]):
     )
 
 
-reg("Input: embedding", outputs=["x"])
+reg("Input: tensor", outputs=["x"], 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 +71,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="same"):
+    if output_dim == "same":
+        oshape = x.shape
+    else:
+        oshape = tuple(*x.shape[:-1], int(output_dim))
+    return Layer(torch.nn.Linear(x.shape, oshape), shape=oshape)
+
+
+class ActivationTypes(enum.Enum):
+    ReLU = "ReLU"
+    Leaky_ReLU = "Leaky ReLU"
+    Tanh = "Tanh"
+    Mish = "Mish"
+
+
+@op("Activation")
+def activation(x, *, type: ActivationTypes = ActivationTypes.ReLU):
+    f = getattr(torch.nn.functional, type.name.lower().replace(" ", "_"))
+    return Layer(f, shape=x.shape)
+
+
 reg("Softmax", inputs=["x"])
 reg(
     "Graph conv",
@@ -70,11 +107,6 @@ 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"])
@@ -128,6 +160,28 @@ 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 OpInput:
+    """Ops get their inputs like this. They have to return a Layer made for this input."""
+
+    id: str
+    shape: tuple[int, ...]
+
+
+@dataclasses.dataclass
+class Layer:
+    """Return this from an op. Must include a module and the shapes of the outputs."""
+
+    module: torch.nn.Module
+    shapes: list[tuple[int, ...]] | None = None  # One for each output.
+    shape: dataclasses.InitVar[tuple[int, ...] | None] = None  # Convenience for single output.
+
+    def __post_init__(self, shape):
+        assert not self.shapes or not shape, "Cannot set both shapes and shape."
+        if shape:
+            self.shapes = [shape]
+
+
 class ColumnSpec(pydantic.BaseModel):
     df: str
     column: str
@@ -306,15 +360,6 @@ def build_model(ws: workspace.Workspace, inputs: dict[str, torch.Tensor]) -> Mod
         outputs = types.SimpleNamespace(**outputs)
         ls = loss_layers if "loss" in regions[node_id] 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(
                     (
@@ -335,6 +380,25 @@ def build_model(ws: workspace.Workspace, inputs: dict[str, torch.Tensor]) -> Mod
                         r = regions.get(n, set())
                         if ("repeat", repeat_id) in r:
                             print(f"repeating {n}")
+            case "Optimizer" | "Input: tensor" | "Input: graph edges" | "Input: sequential":
+                pass
+            case _:
+                op_inputs = []
+                for i in op.inputs.keys():
+                    id = getattr(inputs, i)
+                    op_inputs.append(OpInput(id, shape=sizes.get(id, 1)))
+                if op.func != ops.no_op:
+                    layer = op.func(*op_inputs, **p)
+                else:
+                    layer = Layer(torch.nn.Identity(), shapes=[i.shape for i in op_inputs])
+                input_ids = ", ".join(i.id for i in op_inputs)
+                output_ids = []
+                for o, shape in zip(op.outputs.keys(), layer.shapes):
+                    id = getattr(outputs, o)
+                    output_ids.append(id)
+                    sizes[id] = shape
+                output_ids = ", ".join(output_ids)
+                ls.append((layer.module, f"{input_ids} -> {output_ids}"))
     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)