Redo PyTorch boxes as a separate environment.

server/ops.py

@@ -181,3 +181,8 @@ def op_registration(env: str):
 
 def passive_op_registration(env: str):
   return functools.partial(register_passive_op, env)
+
+def register_area(env, name, params=[]):
+  '''A node that represents an area. It can contain other nodes, but does not restrict movement in any way.'''
+  op = Op(func=no_op, name=name, params={p.name: p for p in params}, inputs={}, outputs={}, type='area')
+  CATALOGS[env][name] = op

server/pytorch_model_ops.py

@@ -1,72 +1,35 @@
-'''Boxes for defining and using PyTorch models.'''
-from enum import Enum
-import inspect
+'''Boxes for defining PyTorch models.'''
 from . import ops
-
-LAYERS = {}
-
-op = ops.op_registration('LynxKite')
-
-@op("Define PyTorch model", sub_nodes=LAYERS)
-def define_pytorch_model(*, sub_flow):
-  print('sub_flow:', sub_flow)
-  return ops.Bundle(other={'model': str(sub_flow)})
-
-@op("Train PyTorch model")
-def train_pytorch_model(model, graph):
-  # import torch # Lazy import because it's slow.
-  return 'hello ' + str(model)
-
-def register_layer(name):
-  def decorator(func):
-    sig = inspect.signature(func)
-    inputs = {
-      name: ops.Input(name=name, type=param.annotation, position='bottom')
-      for name, param in sig.parameters.items()
-      if param.kind != param.KEYWORD_ONLY}
-    params = {
-      name: ops.Parameter.basic(name, param.default, param.annotation)
-      for name, param in sig.parameters.items()
-      if param.kind == param.KEYWORD_ONLY}
-    outputs = {'x': ops.Output(name='x', type='tensor', position='top')}
-    LAYERS[name] = ops.Op(func=func, name=name, params=params, inputs=inputs, outputs=outputs)
-    return func
-  return decorator
-
-@register_layer('LayerNorm')
-def layernorm(x):
-  return 'LayerNorm'
-
-@register_layer('Dropout')
-def dropout(x, *, p=0.5):
-  return f'Dropout ({p})'
-
-@register_layer('Linear')
-def linear(*, output_dim: int):
-  return f'Linear {output_dim}'
-
-class GraphConv(Enum):
-  GCNConv = 'GCNConv'
-  GATConv = 'GATConv'
-  GATv2Conv = 'GATv2Conv'
-  SAGEConv = 'SAGEConv'
-
-@register_layer('Graph Convolution')
-def graph_convolution(x, edges, *, type: GraphConv):
-  return 'GraphConv'
-
-class Nonlinearity(Enum):
-  Mish = 'Mish'
-  ReLU = 'ReLU'
-  Tanh = 'Tanh'
-
-@register_layer('Nonlinearity')
-def nonlinearity(x, *, type: Nonlinearity):
-  return 'ReLU'
-
-def register_area(name, params=[]):
-  '''A node that represents an area. It can contain other nodes, but does not restrict movement in any way.'''
-  op = ops.Op(func=ops.no_op, name=name, params={p.name: p for p in params}, inputs={}, outputs={}, type='area')
-  LAYERS[name] = op
-
-register_area('Repeat', params=[ops.Parameter.basic('times', 1, int)])
+from .ops import Parameter as P
+
+ENV = 'PyTorch model'
+def reg(name, inputs=[], outputs=None, params=[]):
+  if outputs is None:
+    outputs = inputs
+  return ops.register_passive_op(
+    ENV, name,
+    inputs=[ops.Input(name=name, position='bottom', type='tensor') for name in inputs],
+    outputs=[ops.Output(name=name, position='top', type='tensor') for name in outputs],
+    params=params)
+reg('Input: features', outputs=['x'])
+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('Attention', inputs=['q', 'k', 'v'], outputs=['x'])
+reg('LayerNorm', inputs=['x'])
+reg('Dropout', inputs=['x'], params=[P.basic('p', 0.5)])
+reg('Linear', inputs=['x'], params=[P.basic('output_dim', 'same')])
+reg('Graph conv', inputs=['x', 'edges'], outputs=['x'],
+    params=[P.options('type', ['GCNConv', 'GATConv', 'GATv2Conv', 'SAGEConv'])])
+reg('Activation', inputs=['x'],
+    params=[P.options('type', ['ReLU', 'LeakyReLU', 'Tanh', 'Mish'])])
+reg('Supervised loss', inputs=['x', 'y'], outputs=['loss'])
+reg('Triplet loss', inputs=['x', 'x_pos', 'x_neg'], 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_area(ENV, 'Repeat', params=[ops.Parameter.basic('times', 1, int)])

web/src/NodeWithArea.svelte

@@ -45,10 +45,6 @@
   .area {
     border-radius: 10px;
     border: 3px dashed oklch(75% 0.2 55);
-    min-width: 400px;
-    min-height: 400px;
-    max-width: 800px;
-    max-height: 800px;
     z-index: 0 !important;
   }
   .title {