Split build_model() so we can repeat a part.

examples/Model definition

@@ -1,12 +1,5 @@
 {
   "edges": [
-    {
-      "id": "MSE loss 1 Optimizer 2",
-      "source": "MSE loss 1",
-      "sourceHandle": "loss",
-      "target": "Optimizer 2",
-      "targetHandle": "loss"
-    },
     {
       "id": "Repeat 1 Linear 2",
       "source": "Repeat 1",
@@ -14,13 +7,6 @@
       "target": "Linear 2",
       "targetHandle": "x"
     },
-    {
-      "id": "Activation 1 MSE loss 1",
-      "source": "Activation 1",
-      "sourceHandle": "output",
-      "target": "MSE loss 1",
-      "targetHandle": "x"
-    },
     {
       "id": "Linear 2 Activation 1",
       "source": "Linear 2",
@@ -35,72 +21,37 @@
       "target": "Repeat 1",
       "targetHandle": "input"
     },
-    {
-      "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"
+    },
+    {
+      "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: tensor 3 MSE loss 2",
+      "source": "Input: tensor 3",
+      "sourceHandle": "x",
+      "target": "MSE loss 2",
+      "targetHandle": "y"
     }
   ],
   "env": "PyTorch model",
   "nodes": [
-    {
-      "data": {
-        "display": null,
-        "error": null,
-        "input_metadata": null,
-        "meta": {
-          "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"
-              }
-            }
-          },
-          "params": {},
-          "type": "basic"
-        },
-        "params": {},
-        "status": "planned",
-        "title": "MSE loss"
-      },
-      "dragHandle": ".bg-primary",
-      "height": 200.0,
-      "id": "MSE loss 1",
-      "position": {
-        "x": 315.0,
-        "y": -510.0
-      },
-      "type": "basic",
-      "width": 200.0
-    },
     {
       "data": {
         "__execution_delay": 0.0,
@@ -384,7 +335,7 @@
       "height": 200.0,
       "id": "Input: tensor 1",
       "position": {
-        "x": 97.60681762952905,
+        "x": 85.83561484252238,
         "y": 293.6278596776366
       },
       "type": "basic",
@@ -434,8 +385,61 @@
       "height": 200.0,
       "id": "Input: tensor 3",
       "position": {
-        "x": 862.4359094222825,
-        "y": -290.0677203273021
+        "x": 485.8840220312055,
+        "y": -149.86223034126274
+      },
+      "type": "basic",
+      "width": 200.0
+    },
+    {
+      "data": {
+        "display": null,
+        "error": null,
+        "input_metadata": null,
+        "meta": {
+          "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": "None"
+              }
+            }
+          },
+          "params": {},
+          "position": {
+            "x": 937.0,
+            "y": 270.0
+          },
+          "type": "basic"
+        },
+        "params": {},
+        "status": "planned",
+        "title": "MSE loss"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 200.0,
+      "id": "MSE loss 2",
+      "position": {
+        "x": 690.0,
+        "y": -480.0
       },
       "type": "basic",
       "width": 200.0

examples/Model use

@@ -579,54 +579,54 @@
               ],
               "data": [
                 [
-                  "[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.33480108 0.59181517 0.76198453 0.98062384]",
+                  "[1.33480108 1.59181523 1.76198459 1.98062384]",
+                  "[1.3419755697250366, 1.5946478843688965, 1.7717586755752563, 1.9897401332855225]"
                 ],
                 [
-                  "[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.91730917 0.22574073 0.09591609 0.33056474]",
+                  "[1.91730917 1.22574067 1.09591603 1.33056474]",
+                  "[1.900892972946167, 1.2247941493988037, 1.0862866640090942, 1.323314905166626]"
                 ],
                 [
-                  "[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.32565445 0.90939188 0.07488042 0.13730896]",
+                  "[1.32565451 1.90939188 1.07488036 1.13730896]",
+                  "[1.3460955619812012, 1.8960161209106445, 1.0530263185501099, 1.1075329780578613]"
                 ],
                 [
-                  "[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.87608397 0.93200487 0.80169648 0.37758952]",
+                  "[1.87608397 1.93200493 1.80169654 1.37758946]",
+                  "[1.87070894241333, 1.9386992454528809, 1.8151044845581055, 1.3952441215515137]"
                 ],
                 [
-                  "[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.39147133 0.29854035 0.84663737 0.58175623]",
+                  "[1.39147139 1.29854035 1.84663737 1.58175623]",
+                  "[1.3877646923065186, 1.2995290756225586, 1.847062587738037, 1.583693265914917]"
                 ],
                 [
-                  "[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.48507756 0.80808765 0.77162558 0.47834778]",
+                  "[1.48507762 1.80808759 1.77162552 1.47834778]",
+                  "[1.490919828414917, 1.8087174892425537, 1.7757861614227295, 1.4824031591415405]"
                 ],
                 [
-                  "[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.75292218 0.81470108 0.49657214 0.56217098]",
+                  "[1.75292218 1.81470108 1.49657214 1.56217098]",
+                  "[1.7527031898498535, 1.8176040649414062, 1.503413438796997, 1.570152759552002]"
                 ],
                 [
-                  "[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.11693293 0.49860179 0.55020827 0.88832849]",
+                  "[1.11693287 1.49860179 1.55020833 1.88832855]",
+                  "[1.1314976215362549, 1.4944026470184326, 1.546830177307129, 1.8803892135620117]"
                 ],
                 [
-                  "[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.19409031 0.68692201 0.60667384 0.57829887]",
+                  "[1.19409037 1.68692207 1.60667384 1.57829881]",
+                  "[1.2091591358184814, 1.6816589832305908, 1.6011345386505127, 1.5684995651245117]"
                 ],
                 [
-                  "[0.59492421 0.90274489 0.38069052 0.46101224]",
-                  "[1.59492421 1.90274489 1.38069057 1.46101224]",
-                  "[1.6593225002288818, -0.006088308058679104, 1.4240546226501465, 1.570335865020752]"
+                  "[0.62569475 0.9881897  0.83639616 0.9828859 ]",
+                  "[1.62569475 1.9881897  1.83639622 1.98288584]",
+                  "[1.6314740180969238, 1.996805191040039, 1.8592857122421265, 2.0075552463531494]"
                 ]
               ]
             },
@@ -648,10 +648,6 @@
                   "[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]"
-                ],
                 [
                   "[0.76807946 0.98855817 0.08259124 0.01730657]",
                   "[1.76807952 1.98855817 1.0825913  1.01730657]"
@@ -684,10 +680,6 @@
                   "[0.98324287 0.99464184 0.14008355 0.47651017]",
                   "[1.98324287 1.99464178 1.14008355 1.47651017]"
                 ],
-                [
-                  "[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]"
@@ -716,10 +708,6 @@
                   "[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]"
@@ -732,14 +720,14 @@
                   "[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]"
-                ],
                 [
                   "[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]"
@@ -785,20 +773,20 @@
                   "[1.78956437 1.87284744 1.06880784 1.03455889]"
                 ],
                 [
-                  "[0.44330525 0.09997386 0.89025736 0.90507984]",
-                  "[1.44330525 1.09997392 1.89025736 1.90507984]"
+                  "[0.94221359 0.57740951 0.98649532 0.40934443]",
+                  "[1.94221354 1.57740951 1.98649526 1.40934443]"
                 ],
                 [
-                  "[0.72290605 0.96945059 0.68354797 0.15270454]",
-                  "[1.72290611 1.96945059 1.68354797 1.15270448]"
+                  "[0.00497234 0.39319336 0.57054168 0.75150961]",
+                  "[1.00497234 1.39319336 1.57054162 1.75150967]"
                 ],
                 [
-                  "[0.75292218 0.81470108 0.49657214 0.56217098]",
-                  "[1.75292218 1.81470108 1.49657214 1.56217098]"
+                  "[0.44330525 0.09997386 0.89025736 0.90507984]",
+                  "[1.44330525 1.09997392 1.89025736 1.90507984]"
                 ],
                 [
-                  "[0.33480108 0.59181517 0.76198453 0.98062384]",
-                  "[1.33480108 1.59181523 1.76198459 1.98062384]"
+                  "[0.72290605 0.96945059 0.68354797 0.15270454]",
+                  "[1.72290611 1.96945059 1.68354797 1.15270448]"
                 ],
                 [
                   "[0.52784437 0.54268694 0.12358981 0.72116476]",
@@ -808,6 +796,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,13 +828,17 @@
                   "[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.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 ]",
@@ -873,8 +869,12 @@
                   "[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]"
+                ],
+                [
+                  "[0.02235305 0.52774918 0.7331115  0.84358269]",
+                  "[1.02235305 1.52774918 1.7331115  1.84358263]"
                 ],
                 [
                   "[0.6080932  0.56563014 0.32107437 0.72599429]",
@@ -908,6 +908,10 @@
                   "[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]"
@@ -916,10 +920,6 @@
                   "[0.12024075 0.21342516 0.56858408 0.58644271]",
                   "[1.12024069 1.21342516 1.56858408 1.58644271]"
                 ],
-                [
-                  "[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]"
@@ -960,6 +960,10 @@
                   "[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]"
@@ -988,10 +992,6 @@
                   "[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]"
-                ],
                 [
                   "[0.68891573 0.25576538 0.96339929 0.503833  ]",
                   "[1.68891573 1.25576544 1.96339929 1.50383306]"
@@ -1000,7 +1000,7 @@
             }
           },
           "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_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)"
+            "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 0x759ed4f2c360>: 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 0x759ed4f2de40>: END_Repeat_1_output, Input__tensor_3_x -> MSE_loss_2_output\n  (1) - Identity(): MSE_loss_2_output -> 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": []
         },
@@ -1270,7 +1270,7 @@
           "type": "basic"
         },
         "params": {
-          "epochs": "150",
+          "epochs": "1500",
           "input_mapping": "{\"map\":{\"Input__tensor_1_x\":{\"df\":\"df_train\",\"column\":\"x\"},\"Input__tensor_3_x\":{\"df\":\"df_train\",\"column\":\"y\"}}}",
           "model_name": "model"
         },

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

@@ -3,7 +3,6 @@
 import copy
 import enum
 import graphlib
-import types
 
 import pydantic
 from lynxkite.core import ops, workspace
@@ -100,6 +99,11 @@ def activation(x, *, type: ActivationTypes = ActivationTypes.ReLU):
     return Layer(f, shape=x.shape)
 
 
+@op("MSE loss")
+def mse_loss(x, y):
+    return Layer(torch.nn.functional.mse_loss, shape=[1])
+
+
 reg("Softmax", inputs=["x"])
 reg(
     "Graph conv",
@@ -111,7 +115,6 @@ 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(
@@ -161,10 +164,10 @@ def _to_id(*strings: str) -> str:
 
 
 @dataclasses.dataclass
-class OpInput:
+class TensorRef:
     """Ops get their inputs like this. They have to return a Layer made for this input."""
 
-    id: str
+    _id: str
     shape: tuple[int, ...]
 
 
@@ -175,12 +178,21 @@ class Layer:
     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.
+    # Set by ModelBuilder.
+    _origin_id: str | None = None
+    _inputs: list[TensorRef] | None = None
+    _outputs: list[TensorRef] | None = None
 
     def __post_init__(self, shape):
         assert not self.shapes or not shape, "Cannot set both shapes and shape."
         if shape:
             self.shapes = [shape]
 
+    def _for_sequential(self):
+        inputs = ", ".join(i._id for i in self._inputs)
+        outputs = ", ".join(o._id for o in self._outputs)
+        return self.module, f"{inputs} -> {outputs}"
+
 
 class ColumnSpec(pydantic.BaseModel):
     df: str
@@ -246,204 +258,212 @@ class ModelConfig:
         }
 
 
-def _add_op(op, params, inputs, outputs, sizes, layers):
-    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, **params)
-    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)
-    layers.append((layer.module, f"{input_ids} -> {output_ids}"))
-
-
-def _all_dependencies(node: str, dependencies: dict[str, list[str]]) -> set[str]:
-    """Returns all dependencies of a node."""
-    deps = set()
-    for dep in dependencies[node]:
-        deps.add(dep)
-        deps.update(_all_dependencies(dep, dependencies))
-    return deps
-
-
 def build_model(ws: workspace.Workspace, inputs: dict[str, torch.Tensor]) -> 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}
-    inv_dependencies = {n.id: [] for n in ws.nodes}
-    in_edges = {}
-    out_edges = {}
-    repeats = []
-    for e in ws.edges:
-        if nodes[e.target].data.title == "Repeat":
-            repeats.append(e.target)
-        dependencies[e.target].append(e.source)
-        inv_dependencies[e.source].append(e.target)
-        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)
-        )
-    # Split repeat boxes into start and end, and insert them into the flow.
-    # TODO: Think about recursive repeats.
-    for repeat in repeats:
-        start_id = f"START {repeat}"
-        end_id = f"END {repeat}"
-        # repeat -> first <- real_input
-        # ...becomes...
-        # real_input -> start -> first
-        first, firsth = out_edges[repeat]["output"][0]
-        [(real_input, real_inputh)] = [
-            k for k in in_edges[first][firsth] if k != (repeat, "output")
-        ]
-        dependencies[first].remove(repeat)
-        dependencies[first].append(start_id)
-        dependencies[start_id] = [real_input]
-        out_edges[real_input][real_inputh] = [
-            k if k != (first, firsth) else (start_id, "input")
-            for k in out_edges[real_input][real_inputh]
-        ]
-        in_edges[start_id] = {"input": [(real_input, real_inputh)]}
-        out_edges[start_id] = {"output": [(first, firsth)]}
-        in_edges[first][firsth] = [(start_id, "output")]
-        # repeat <- last -> real_output
-        # ...becomes...
-        # last -> end -> real_output
-        last, lasth = in_edges[repeat]["input"][0]
-        [(real_output, real_outputh)] = [
-            k for k in out_edges[last][lasth] if k != (repeat, "input")
-        ]
-        del dependencies[repeat]
-        dependencies[end_id] = [last]
-        dependencies[real_output].append(end_id)
-        out_edges[last][lasth] = [(end_id, "input")]
-        in_edges[end_id] = {"input": [(last, lasth)]}
-        out_edges[end_id] = {"output": [(real_output, real_outputh)]}
-        in_edges[real_output][real_outputh] = [
-            k if k != (last, lasth) else (end_id, "output")
-            for k in in_edges[real_output][real_outputh]
-        ]
-    # Walk the graph in topological order.
-    sizes = {}
-    for k, i in inputs.items():
-        sizes[k] = i.shape[-1]
-    ts = graphlib.TopologicalSorter(dependencies)
-    layers = []
-    loss_layers = []
-    regions: dict[str, set[str]] = {node_id: set() for node_id in dependencies}
-    cfg = {}
-    used_in_model = set()
-    made_in_model = set()
-    used_in_loss = set()
-    made_in_loss = set()
-    for node_id in ts.static_order():
+    builder = ModelBuilder(ws, inputs)
+    return builder.build_model()
+
+
+class ModelBuilder:
+    """The state shared between methods that are used to build the model."""
+
+    def __init__(self, ws: workspace.Workspace, inputs: dict[str, torch.Tensor]):
+        self.catalog = ops.CATALOGS[ENV]
+        optimizers = []
+        self.nodes = {}
+        for node in ws.nodes:
+            self.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}"
+        [self.optimizer] = optimizers
+        self.dependencies = {n.id: [] for n in ws.nodes}
+        self.in_edges = {}
+        self.out_edges = {}
+        repeats = []
+        for e in ws.edges:
+            if self.nodes[e.target].data.title == "Repeat":
+                repeats.append(e.target)
+            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:
+            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.dependencies}
+        for k, v in self.dependencies.items():
+            for i in v:
+                self.inv_dependencies[i].append(k)
+        self.sizes = {}
+        for k, i in inputs.items():
+            self.sizes[k] = i.shape[-1]
+        self.layers = []
+
+    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."""
         if node_id.startswith("START "):
-            node = nodes[node_id.removeprefix("START ")]
+            node = self.nodes[node_id.removeprefix("START ")]
         elif node_id.startswith("END "):
-            node = nodes[node_id.removeprefix("END ")]
+            node = self.nodes[node_id.removeprefix("END ")]
         else:
-            node = nodes[node_id]
+            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]:
-            regions[node_id] |= regions[b]
-        if "loss" in t:
-            regions[node_id].add("loss")
         inputs = {}
-        for n in in_edges.get(node_id, []):
-            for b, h in in_edges[node_id][n]:
+        for n in self.in_edges.get(node_id, []):
+            for b, h in self.in_edges[node_id][n]:
                 i = _to_id(b, h)
                 inputs[n] = i
-                if "loss" in regions[node_id]:
-                    used_in_loss.add(i)
-                else:
-                    used_in_model.add(i)
         outputs = {}
-        for out in out_edges.get(node_id, []):
+        for out in self.out_edges.get(node_id, []):
             i = _to_id(node_id, out)
             outputs[out] = i
-            if not t.startswith("Input:"):  # The outputs of inputs are not "made" by us.
-                if "loss" in regions[node_id]:
-                    made_in_loss.add(i)
-                else:
-                    made_in_model.add(i)
-        inputs = types.SimpleNamespace(**inputs)
-        outputs = types.SimpleNamespace(**outputs)
-        ls = loss_layers if "loss" in regions[node_id] else layers
         match t:
-            case "MSE loss":
-                ls.append(
-                    (
-                        torch.nn.functional.mse_loss,
-                        f"{inputs.x}, {inputs.y} -> {outputs.loss}",
-                    )
-                )
             case "Repeat":
-                ls.append((torch.nn.Identity(), f"{inputs.input} -> {outputs.output}"))
-                sizes[outputs.output] = sizes.get(inputs.input, 1)
-                if node_id.startswith("START "):
-                    regions[node_id].add(("repeat", node_id.removeprefix("START ")))
-                else:
+                if node_id.startswith("END "):
                     repeat_id = node_id.removeprefix("END ")
                     start_id = f"START {repeat_id}"
                     print(f"repeat {repeat_id} ending")
-                    after_start = _all_dependencies(start_id, inv_dependencies)
-                    after_end = _all_dependencies(node_id, inv_dependencies)
-                    before_end = _all_dependencies(node_id, dependencies)
-                    affected_nodes = after_start - after_end
+                    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}"
                     )
-                    regions[node_id].remove(("repeat", repeat_id))
                     for n in repeated_nodes:
                         print(f"repeating {n}")
             case "Optimizer" | "Input: tensor" | "Input: graph edges" | "Input: sequential":
-                pass
-            case _:
-                _add_op(op, p, inputs, outputs, sizes, ls)
-    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)
+                return
+        layer = self.run_op(op, p, inputs, outputs)
+        layer._origin_id = node_id
+        self.layers.append(layer)
+
+    def run_op(self, op, params, inputs: dict[str, str], outputs: dict[str, str]) -> Layer:
+        """Returns the layer produced by this op."""
+        op_inputs = [
+            TensorRef(inputs[i], shape=self.sizes.get(inputs[i], 1)) for i in op.inputs.keys()
+        ]
+        if op.func != ops.no_op:
+            layer = op.func(*op_inputs, **params)
+        else:
+            layer = Layer(torch.nn.Identity(), shapes=[i.shape for i in op_inputs])
+        layer._inputs = op_inputs
+        layer._outputs = []
+        for o, shape in zip(op.outputs.keys(), layer.shapes):
+            layer._outputs.append(TensorRef(outputs[o], shape=shape))
+            self.sizes[outputs[o]] = shape
+        return layer
+
+    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._id for layer in layers for input in layer._inputs)
+        used_in_loss = set(input._id for layer in loss_layers for input in layer._inputs)
+        made_in_model = set(output._id for layer in layers for output in layer._outputs)
+        made_in_loss = set(output._id 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"]
+        p = op.convert_params(self.nodes[self.optimizer].data.params)
+        o = getattr(torch.optim, p["type"].name)
+        cfg["optimizer"] = o(cfg["model"].parameters(), lr=p["lr"])
+        print(cfg)
+        return ModelConfig(**cfg)
 
 
 def to_tensors(b: core.Bundle, m: ModelMapping | None) -> dict[str, torch.Tensor]: