Split ops from infrastructure code.

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,7 @@ from lynxkite.core import ops
 from collections import deque
 
 from tqdm import tqdm
-from . import core, pytorch_model_ops
+from . import core, pytorch
 from lynxkite.core import workspace
 import grandcypher
 import joblib
@@ -347,7 +347,7 @@ def define_model(
     assert model_workspace, "Model workspace is unset."
     ws = load_ws(model_workspace)
     # Build the model without inputs, to get its interface.
-    m = pytorch_model_ops.build_model(ws)
+    m = pytorch.core.build_model(ws)
     m.source_workspace = model_workspace
     bundle = bundle.copy()
     bundle.other[save_as] = m
@@ -356,15 +356,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 +379,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 +406,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 core  # noqa
+from . import ops  # noqa

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

@@ -1,16 +1,14 @@
 """Boxes 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)

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

@@ -0,0 +1,118 @@
+"""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 .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("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=[],
+)

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 import pytorch
 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,7 +46,7 @@ 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"},
             "lin": {"title": "Linear", "output_dim": 4},
@@ -65,7 +65,7 @@ async def test_build_model():
     )
     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})
     assert loss < 0.1
@@ -77,7 +77,7 @@ async def test_build_model():
 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"},
                 "lin": {"title": "Linear", "output_dim": 8},
@@ -99,17 +99,17 @@ async def test_build_model_with_repeat():
         )
 
     # 1 repetition
-    m = pytorch_model_ops.build_model(repeated_ws(1))
+    m = pytorch.core.build_model(repeated_ws(1))
     assert summarize_layers(m) == "IL<II"
     assert summarize_connections(m) == "e->S S->l l->a a->E E->E"
 
     # 2 repetitions
-    m = pytorch_model_ops.build_model(repeated_ws(2))
+    m = pytorch.core.build_model(repeated_ws(2))
     assert summarize_layers(m) == "IL<IL<II"
     assert summarize_connections(m) == "e->S S->l l->a a->S S->l l->a a->E E->E"
 
     # 3 repetitions
-    m = pytorch_model_ops.build_model(repeated_ws(3))
+    m = pytorch.core.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"