Delay some imports to speed up startup.

lynxkite-graph-analytics/pyproject.toml

@@ -26,6 +26,7 @@ dev = [
     "pytest-asyncio>=0.26.0",
 ]
 gpu = [
+    "cuml-cu12>=25.2.1",
     "nx-cugraph-cu12>=25.4.0",
     "polars[gpu]>=1.25.2",
 ]

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

@@ -8,7 +8,6 @@ from lynxkite.core import workspace
 from .pytorch import pytorch_core
 from lynxkite.core import ops
 from tqdm import tqdm
-import umap
 import joblib
 import pandas as pd
 import pathlib
@@ -178,15 +177,19 @@ def view_vectors(
     min_dist: float = 0.1,
     metric: UMAPMetric = UMAPMetric.euclidean,
 ):
+    try:
+        from cuml.manifold.umap import UMAP
+    except ImportError:
+        from umap import UMAP
     vec = np.stack(bundle.dfs[table_name][vector_column].to_numpy())
-    _umap = functools.partial(
-        umap.UMAP,
+    umap = functools.partial(
+        UMAP,
         n_neighbors=n_neighbors,
         min_dist=min_dist,
         metric=metric.value,
     )
-    proj = _umap(n_components=2).fit_transform(vec)
-    color = _umap(n_components=1).fit_transform(vec)
+    proj = umap(n_components=2).fit_transform(vec)
+    color = umap(n_components=1).fit_transform(vec)
     data = [[*p.tolist(), "", c.item()] for p, c in zip(proj, color)]
     if label_column:
         for i, row in enumerate(bundle.dfs[table_name][label_column]):

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

@@ -6,7 +6,6 @@ import graphlib
 import pydantic
 from lynxkite.core import ops, workspace
 import torch
-import torch_geometric.nn as pyg_nn
 import dataclasses
 from .. import core
 
@@ -319,6 +318,8 @@ class ModelBuilder:
         return self.get_config()
 
     def get_config(self) -> ModelConfig:
+        import torch_geometric.nn as pyg_nn
+
         # Split the design into model and loss.
         model_nodes = set()
         for node_id in self.nodes:

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

@@ -4,7 +4,6 @@ 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 .pytorch_core import op, reg, ENV
 
 reg("Input: tensor", outputs=["output"], params=[P.basic("name")])
@@ -61,6 +60,8 @@ def dropout(x, *, p=0.0):
 
 @op("Linear", weights=True)
 def linear(x, *, output_dim=1024):
+    import torch_geometric.nn as pyg_nn
+
     return pyg_nn.Linear(-1, output_dim)