Rename core/ops to pytorch_core/pytorch_ops.

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

@@ -8,7 +8,8 @@ from lynxkite.core import ops
 from collections import deque
 
 from tqdm import tqdm
-from . import core, pytorch
+from . import core
+from .pytorch import pytorch_core
 from lynxkite.core import workspace
 import grandcypher
 import joblib
@@ -347,7 +348,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.core.build_model(ws)
+    m = pytorch_core.build_model(ws)
     m.source_workspace = model_workspace
     bundle = bundle.copy()
     bundle.other[save_as] = m
@@ -356,15 +357,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.core.ModelMapping):
+class ModelInferenceInputMapping(pytorch_core.ModelMapping):
     pass
 
 
-class ModelTrainingInputMapping(pytorch.core.ModelMapping):
+class ModelTrainingInputMapping(pytorch_core.ModelMapping):
     pass
 
 
-class ModelOutputMapping(pytorch.core.ModelMapping):
+class ModelOutputMapping(pytorch_core.ModelMapping):
     pass
 
 
@@ -379,7 +380,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.core.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 +407,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.core.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

@@ -1,2 +1,2 @@
-from . import core  # noqa
-from . import ops  # noqa
+from . import pytorch_core  # noqa
+from . import pytorch_ops  # noqa

lynxkite-graph-analytics/src/lynxkite_graph_analytics/pytorch/{ops.py → pytorch_ops.py}

@@ -5,7 +5,7 @@ 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
+from .pytorch_core import op, reg, ENV
 
 reg("Input: tensor", outputs=["output"], params=[P.basic("name")])
 reg("Input: graph edges", outputs=["edges"])

lynxkite-graph-analytics/tests/test_pytorch_model_ops.py

@@ -1,5 +1,5 @@
 from lynxkite.core import workspace
-from lynxkite_graph_analytics import pytorch
+from lynxkite_graph_analytics.pytorch import pytorch_core
 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.core.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.core.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.core.ModelConfig) -> str:
 
 async def test_build_model():
     ws = make_ws(
-        pytorch.core.ENV,
+        pytorch_core.ENV,
         {
             "input": {"title": "Input: tensor"},
             "lin": {"title": "Linear", "output_dim": 4},
@@ -67,7 +67,7 @@ async def test_build_model():
     )
     x = torch.rand(100, 4)
     y = x + 1
-    m = pytorch.core.build_model(ws)
+    m = pytorch_core.build_model(ws)
     for i in range(1000):
         loss = m.train({"input_output": x, "label_output": y})
     assert loss < 0.1
@@ -79,7 +79,7 @@ async def test_build_model():
 async def test_build_model_with_repeat():
     def repeated_ws(times):
         return make_ws(
-            pytorch.core.ENV,
+            pytorch_core.ENV,
             {
                 "input": {"title": "Input: tensor"},
                 "lin": {"title": "Linear", "output_dim": 8},
@@ -103,17 +103,17 @@ async def test_build_model_with_repeat():
         )
 
     # 1 repetition
-    m = pytorch.core.build_model(repeated_ws(1))
+    m = pytorch_core.build_model(repeated_ws(1))
     assert summarize_layers(m) == "IL<III"
     assert summarize_connections(m) == "i->S S->l l->a a->E E->o o->o"
 
     # 2 repetitions
-    m = pytorch.core.build_model(repeated_ws(2))
+    m = pytorch_core.build_model(repeated_ws(2))
     assert summarize_layers(m) == "IL<IL<III"
     assert summarize_connections(m) == "i->S S->l l->a a->S S->l l->a a->E E->o o->o"
 
     # 3 repetitions
-    m = pytorch.core.build_model(repeated_ws(3))
+    m = pytorch_core.build_model(repeated_ws(3))
     assert summarize_layers(m) == "IL<IL<IL<III"
     assert summarize_connections(m) == "i->S S->l l->a a->S S->l l->a a->S S->l l->a a->E E->o o->o"