Merge pull request #184 from biggraph/darabos-ode-gnn-v2

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

@@ -54,6 +54,7 @@ function getHandles(inputs: any[], outputs: any[]) {
 }
 
 const OP_COLORS: { [key: string]: string } = {
+  gray: "oklch(95% 0 0)",
   pink: "oklch(75% 0.2 0)",
   orange: "oklch(75% 0.2 55)",
   green: "oklch(75% 0.2 150)",

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

@@ -240,7 +240,7 @@ def op(
     """Decorator for defining an operation."""
 
     def decorator(func):
-        doc = get_doc(func)
+        doc = parse_doc(func.__doc__)
         sig = inspect.signature(func)
         _view = view
         if view == "matplotlib":
@@ -436,16 +436,17 @@ def run_user_script(script_path: pathlib.Path):
     spec.loader.exec_module(module)
 
 
-def get_doc(func):
+@functools.cache
+def parse_doc(doc):
     """Griffe is an optional dependency. When available, we returned the parsed docstring."""
     try:
         import griffe
     except ImportError:
-        return func.__doc__
-    if func.__doc__ is None:
+        return doc
+    if doc is None:
         return None
-    if "----" in func.__doc__:
-        doc = griffe.Docstring(func.__doc__).parse("numpy")
+    if "----" in doc:
+        doc = griffe.Docstring(doc).parse("numpy")
     else:
-        doc = griffe.Docstring(func.__doc__).parse("google")
+        doc = griffe.Docstring(doc).parse("google")
     return json.loads(json.dumps(doc, cls=griffe.JSONEncoder))

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

@@ -6,24 +6,23 @@ from lynxkite.core.ops import Parameter as P
 import torch
 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"], params=[P.basic("name")])
-reg("Output", inputs=["x"], outputs=["x"], params=[P.basic("name")])
+reg("Input: tensor", outputs=["output"], params=[P.basic("name")], color="gray")
+reg("Input: graph edges", outputs=["edges"], params=[P.basic("name")], color="gray")
+reg("Input: sequential", outputs=["y"], params=[P.basic("name")], color="gray")
+reg("Output", inputs=["x"], outputs=["x"], params=[P.basic("name")], color="gray")
 
 
 @op("LSTM", weights=True)
 def lstm(x, *, input_size=1024, hidden_size=1024, dropout=0.0):
-    return torch.nn.LSTM(input_size, hidden_size, dropout=0.0)
+    return torch.nn.LSTM(input_size, hidden_size, dropout=dropout)
 
 
 reg(
-    "Neural ODE",
+    "Neural ODE with MLP",
     color="blue",
-    inputs=["x"],
+    inputs=["x", "y0", "t"],
+    outputs=["y"],
     params=[
-        P.basic("relative_tolerance"),
-        P.basic("absolute_tolerance"),
         P.options(
             "method",
             [
@@ -39,6 +38,11 @@ reg(
                 "implicit_adams",
             ],
         ),
+        P.basic("relative_tolerance"),
+        P.basic("absolute_tolerance"),
+        P.basic("mlp_layers"),
+        P.basic("mlp_hidden_size"),
+        P.options("mlp_activation", ["ReLU", "Tanh", "Sigmoid"]),
     ],
 )
 
@@ -66,6 +70,13 @@ def linear(x, *, output_dim=1024):
     return pyg_nn.Linear(-1, output_dim)
 
 
+@op("Mean pool")
+def mean_pool(x):
+    import torch_geometric.nn as pyg_nn
+
+    return pyg_nn.global_mean_pool
+
+
 class ActivationTypes(str, enum.Enum):
     ReLU = "ReLU"
     Leaky_ReLU = "Leaky ReLU"
@@ -93,11 +104,39 @@ def softmax(x, *, dim=1):
     return torch.nn.Softmax(dim=dim)
 
 
+@op("Embedding", weights=True)
+def embedding(x, *, num_embeddings: int, embedding_dim: int):
+    return torch.nn.Embedding(num_embeddings, embedding_dim)
+
+
 @op("Concatenate")
 def concatenate(a, b):
     return lambda a, b: torch.concatenate(*torch.broadcast_tensors(a, b))
 
 
+reg(
+    "Pick element by index",
+    inputs=["x", "index"],
+    outputs=["x_i"],
+)
+reg(
+    "Pick element by constant",
+    inputs=["x"],
+    outputs=["x_i"],
+    params=[ops.Parameter.basic("index", "0")],
+)
+reg(
+    "Take first n",
+    inputs=["x"],
+    outputs=["x"],
+    params=[ops.Parameter.basic("n", 1, int)],
+)
+reg(
+    "Drop first n",
+    inputs=["x"],
+    outputs=["x"],
+    params=[ops.Parameter.basic("n", 1, int)],
+)
 reg(
     "Graph conv",
     color="blue",
@@ -105,6 +144,15 @@ reg(
     outputs=["x"],
     params=[P.options("type", ["GCNConv", "GATConv", "GATv2Conv", "SAGEConv"])],
 )
+reg(
+    "Heterogeneous graph conv",
+    inputs=["node_embeddings", "edge_modules"],
+    outputs=["x"],
+    params=[
+        ops.Parameter.basic("node_embeddings_order"),
+        ops.Parameter.basic("edge_modules_order"),
+    ],
+)
 
 reg("Triplet margin loss", inputs=["x", "x_pos", "x_neg"], outputs=["loss"])
 reg("Cross-entropy loss", inputs=["x", "y"], outputs=["loss"])
@@ -125,7 +173,7 @@ reg(
                 "Galore AdamW",
             ],
         ),
-        P.basic("lr", 0.001),
+        P.basic("lr", 0.0001),
     ],
     color="green",
 )