modssc 0.2.1.dev0__tar.gz → 0.2.2__tar.gz

{modssc-0.2.1.dev0 → modssc-0.2.2}/CHANGELOG.md

@@ -4,9 +4,17 @@ All notable changes to this project will be documented in this file.
 
 The format is based on "Keep a Changelog", and this project adheres to Semantic Versioning.
 
+## 0.2.2
+- Ajout du support de `activation` pour GraphSAGE inductif.
+- Ajout de `core.to_torch` dans les configs TriNet vision (best + smoke).
+- Mise a jour des tests GraphSAGE.
+
+## 0.2.1
+- Bump version metadata.
+
 ## 0.2.0
-- Amelioration des performances des pipelines inductifs et des chemins critiques.
-- Couverture de tests renforcee pour securiser les nouvelles optimisations.
+- Improved inductive pipeline performance and critical paths.
+- Strengthened test coverage to secure the new optimizations.
 
 ## 0.1.2
 - Fix misc bugs

{modssc-0.2.1.dev0 → modssc-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: modssc
-Version: 0.2.1.dev0
+Version: 0.2.2
 Summary: ModSSC: a modular framework for semi-supervised classification on heterogeneous data
 Project-URL: Homepage, https://github.com/ModSSC/ModSSC
 Project-URL: Repository, https://github.com/ModSSC/ModSSC

{modssc-0.2.1.dev0 → modssc-0.2.2}/docs/index.md

@@ -33,7 +33,7 @@ ModSSC is a modular framework for semi-supervised classification across heteroge
 - [Transductive tutorial](tutorials/transductive-toy.md)
 
 ## Version
-Current version is `0.2.0`, sourced from [`src/modssc/__about__.py`](https://github.com/ModSSC/ModSSC/blob/main/src/modssc/__about__.py) and referenced by the Hatch version config in [`pyproject.toml`](https://github.com/ModSSC/ModSSC/blob/main/pyproject.toml). <sup class="cite"><a href="#source-18">[18]</a><a href="#source-2">[2]</a></sup>
+Current version is `0.2.2`, sourced from [`src/modssc/__about__.py`](https://github.com/ModSSC/ModSSC/blob/main/src/modssc/__about__.py) and referenced by the Hatch version config in [`pyproject.toml`](https://github.com/ModSSC/ModSSC/blob/main/pyproject.toml). <sup class="cite"><a href="#source-18">[18]</a><a href="#source-2">[2]</a></sup>
 
 
 ## Project status and support

{modssc-0.2.1.dev0 → modssc-0.2.2}/src/modssc/__about__.py

@@ -1,4 +1,4 @@
 # SPDX-FileCopyrightText: 2024-present ModSSC contributors
 #
 # SPDX-License-Identifier: MIT
-__version__ = "0.2.1.dev0"
+__version__ = "0.2.2"

{modssc-0.2.1.dev0 → modssc-0.2.2}/src/modssc/inductive/adapters/torch.py

@@ -4,6 +4,8 @@ from collections.abc import Mapping
 from dataclasses import dataclass
 from typing import Any
 
+import numpy as np
+
 from ..backends import torch_backend
 from ..base import InductiveDatasetLike
 from ..errors import InductiveValidationError
@@ -177,6 +179,32 @@ def to_torch_dataset(
                     f"Tensor device mismatch: expected {expected}, got {t.device}"
                 )
 
+    meta = data.meta
+    if isinstance(meta, Mapping):
+        meta = dict(meta)
+        torch = _torch()
+
+        def _device_of(obj: Any) -> Any | None:
+            if obj is None:
+                return None
+            if isinstance(obj, dict) and "x" in obj:
+                return obj["x"].device
+            return getattr(obj, "device", None)
+
+        idx_device = _device_of(X_u) or _device_of(X_u_w) or _device_of(X_u_s) or _device_of(X_l)
+        if idx_device is not None:
+            for key in ("idx_u", "unlabeled_idx", "unlabeled_indices"):
+                if key not in meta:
+                    continue
+                idx_val = meta[key]
+                if isinstance(idx_val, torch.Tensor):
+                    if idx_val.dtype != torch.int64 or idx_val.device != idx_device:
+                        meta[key] = idx_val.to(device=idx_device, dtype=torch.int64)
+                else:
+                    meta[key] = torch.as_tensor(
+                        np.asarray(idx_val), device=idx_device, dtype=torch.int64
+                    )
+
     return TorchDataset(
         X_l=X_l,
         y_l=y_l,
@@ -184,5 +212,5 @@ def to_torch_dataset(
         X_u_w=X_u_w,
         X_u_s=X_u_s,
         views=views,
-        meta=data.meta,
+        meta=meta,
     )

{modssc-0.2.1.dev0 → modssc-0.2.2}/src/modssc/inductive/deep/bundles.py

@@ -670,8 +670,23 @@ def _build_graphsage_bundle(
     except ImportError as e:
         raise ImportError("torch_geometric is required for graphsage_inductive") from e
 
+    hidden_sizes = params.get("hidden_sizes")
     hidden_channels = int(params.get("hidden_channels", 128))
     num_layers = int(params.get("num_layers", 2))
+    if hidden_sizes is not None:
+        if isinstance(hidden_sizes, int):
+            hidden_sizes = [int(hidden_sizes)]
+        elif isinstance(hidden_sizes, (list, tuple)):
+            hidden_sizes = [int(h) for h in hidden_sizes]
+        else:
+            raise InductiveValidationError("hidden_sizes must be an int or a sequence of ints.")
+        if any(h <= 0 for h in hidden_sizes):
+            raise InductiveValidationError("hidden_sizes must be positive.")
+        if "num_layers" in params and num_layers != len(hidden_sizes) + 1:
+            raise InductiveValidationError(
+                "num_layers must equal len(hidden_sizes) + 1 when hidden_sizes is provided."
+            )
+        num_layers = len(hidden_sizes) + 1
     dropout = float(params.get("dropout", 0.5))
     lr = float(params.get("lr", 1e-2))
     weight_decay = float(params.get("weight_decay", 5e-4))
@@ -690,20 +705,11 @@ def _build_graphsage_bundle(
         )
 
     class _GraphSAGEWrapper(torch.nn.Module):
-        def __init__(
-            self,
-            in_channels: int,
-            hidden_channels: int,
-            num_layers: int,
-            out_channels: int,
-            dropout: float,
-        ):
+        def __init__(self, layer_sizes: list[int], dropout: float):
             super().__init__()
             self.convs = torch.nn.ModuleList()
-            self.convs.append(SAGEConv(in_channels, hidden_channels))
-            for _ in range(num_layers - 2):
-                self.convs.append(SAGEConv(hidden_channels, hidden_channels))
-            self.convs.append(SAGEConv(hidden_channels, out_channels))
+            for in_channels, out_channels in zip(layer_sizes[:-1], layer_sizes[1:], strict=False):
+                self.convs.append(SAGEConv(in_channels, out_channels))
             self.dropout = dropout
 
         def forward(self, x: Any):
@@ -726,13 +732,11 @@ def _build_graphsage_bundle(
             return {"logits": logits, "feat": feat}
 
     torch.manual_seed(int(seed))
-    model = _GraphSAGEWrapper(
-        in_channels=int(in_channels),
-        hidden_channels=hidden_channels,
-        num_layers=num_layers,
-        out_channels=int(num_classes),
-        dropout=dropout,
-    ).to(device)
+    if hidden_sizes is not None:
+        layer_sizes = [int(in_channels), *hidden_sizes, int(num_classes)]
+    else:
+        layer_sizes = [int(in_channels)] + [hidden_channels] * (num_layers - 1) + [int(num_classes)]
+    model = _GraphSAGEWrapper(layer_sizes=layer_sizes, dropout=dropout).to(device)
 
     optimizer = torch.optim.AdamW(
         model.parameters(), lr=float(lr), weight_decay=float(weight_decay)

{modssc-0.2.1.dev0 → modssc-0.2.2}/src/modssc/inductive/methods/adamatch.py

@@ -9,6 +9,7 @@ from modssc.inductive.base import InductiveMethod, MethodInfo
 from modssc.inductive.deep import TorchModelBundle
 from modssc.inductive.errors import InductiveValidationError
 from modssc.inductive.methods.deep_utils import (
+    concat_data,
     cycle_batch_indices,
     cycle_batches,
     ensure_float_tensor,
@@ -247,13 +248,13 @@ class AdaMatchMethod(InductiveMethod):
                 x_us = _slice_dict(X_u_s, idx_u, int(_get_len(X_u_s)))
 
                 if bool(self.spec.use_cat):
-                    inputs = torch.cat([x_lb, x_uw, x_us], dim=0)
+                    inputs = concat_data([x_lb, x_uw, x_us])
                     logits = extract_logits(model(inputs))
                     if int(logits.ndim) != 2:
                         raise InductiveValidationError("Model logits must be 2D (batch, classes).")
-                    num_lb = int(x_lb.shape[0])
-                    num_u = int(x_uw.shape[0])
-                    expected = num_lb + num_u + int(x_us.shape[0])
+                    num_lb = int(_get_len(x_lb))
+                    num_u = int(_get_len(x_uw))
+                    expected = num_lb + num_u + int(_get_len(x_us))
                     if int(logits.shape[0]) != expected:
                         raise InductiveValidationError(
                             "Concatenated logits batch size does not match inputs."

{modssc-0.2.1.dev0 → modssc-0.2.2}/src/modssc/inductive/methods/adsh.py

@@ -10,13 +10,17 @@ from modssc.inductive.base import InductiveMethod, MethodInfo
 from modssc.inductive.deep import TorchModelBundle
 from modssc.inductive.errors import InductiveValidationError
 from modssc.inductive.methods.deep_utils import (
+    concat_data,
     cycle_batch_indices,
     cycle_batches,
     ensure_float_tensor,
     ensure_model_bundle,
     ensure_model_device,
     extract_logits,
+    get_torch_device,
+    get_torch_len,
     num_batches,
+    slice_data,
 )
 from modssc.inductive.methods.utils import (
     detect_backend,
@@ -43,7 +47,7 @@ def _update_scores(
     torch = optional_import("torch", extra="inductive-torch")
     if score.ndim != 1:
         raise InductiveValidationError("ADSH score must be a 1D tensor.")
-    if score.device != X_u_w.device:
+    if score.device != get_torch_device(X_u_w):
         raise InductiveValidationError("ADSH score must be on the same device as X_u_w.")
 
     n_classes = int(score.shape[0])
@@ -51,8 +55,8 @@ def _update_scores(
     was_training = model.training
     model.eval()
     with torch.no_grad():
-        for start in range(0, int(X_u_w.shape[0]), int(batch_size)):
-            batch = X_u_w[start : start + int(batch_size)]
+        for start in range(0, int(get_torch_len(X_u_w)), int(batch_size)):
+            batch = slice_data(X_u_w, slice(start, start + int(batch_size)))
             logits = extract_logits(model(batch))
             if int(logits.ndim) != 2:
                 raise InductiveValidationError("Model logits must be 2D (batch, classes).")
@@ -184,15 +188,15 @@ class ADSHMethod(InductiveMethod):
         X_u_s = ds.X_u_s
         logger.info(
             "ADSH sizes: n_labeled=%s n_unlabeled=%s",
-            int(X_l.shape[0]),
-            int(X_u_w.shape[0]),
+            int(get_torch_len(X_l)),
+            int(get_torch_len(X_u_w)),
         )
 
-        if int(X_l.shape[0]) == 0:
+        if int(get_torch_len(X_l)) == 0:
             raise InductiveValidationError("X_l must be non-empty.")
-        if int(X_u_w.shape[0]) == 0 or int(X_u_s.shape[0]) == 0:
+        if int(get_torch_len(X_u_w)) == 0 or int(get_torch_len(X_u_s)) == 0:
             raise InductiveValidationError("X_u_w and X_u_s must be non-empty.")
-        if int(X_u_w.shape[0]) != int(X_u_s.shape[0]):
+        if int(get_torch_len(X_u_w)) != int(get_torch_len(X_u_s)):
             raise InductiveValidationError("X_u_w and X_u_s must have the same number of rows.")
 
         ensure_float_tensor(X_l, name="X_l")
@@ -207,7 +211,7 @@ class ADSHMethod(InductiveMethod):
         bundle = ensure_model_bundle(self.spec.model_bundle)
         model = bundle.model
         optimizer = bundle.optimizer
-        ensure_model_device(model, device=X_l.device)
+        ensure_model_device(model, device=get_torch_device(X_l))
 
         if int(self.spec.batch_size) <= 0:
             raise InductiveValidationError("batch_size must be >= 1.")
@@ -223,7 +227,7 @@ class ADSHMethod(InductiveMethod):
         was_training = model.training
         model.eval()
         with torch.no_grad():
-            init_logits = extract_logits(model(X_l[:1]))
+            init_logits = extract_logits(model(slice_data(X_l, slice(0, 1))))
         if was_training:
             model.train()
         if int(init_logits.ndim) != 2:
@@ -236,12 +240,12 @@ class ADSHMethod(InductiveMethod):
         score = torch.full(
             (n_classes,),
             float(self.spec.p_cutoff),
-            device=X_l.device,
+            device=get_torch_device(X_l),
             dtype=init_logits.dtype,
         )
 
-        steps_l = num_batches(int(X_l.shape[0]), int(self.spec.batch_size))
-        steps_u = num_batches(int(X_u_w.shape[0]), int(self.spec.batch_size))
+        steps_l = num_batches(int(get_torch_len(X_l)), int(self.spec.batch_size))
+        steps_u = num_batches(int(get_torch_len(X_u_w)), int(self.spec.batch_size))
         steps_per_epoch = max(int(steps_l), int(steps_u))
 
         gen_l = torch.Generator().manual_seed(int(seed))
@@ -267,24 +271,24 @@ class ADSHMethod(InductiveMethod):
                 steps=steps_per_epoch,
             )
             iter_u_idx = cycle_batch_indices(
-                int(X_u_w.shape[0]),
+                int(get_torch_len(X_u_w)),
                 batch_size=int(self.spec.batch_size),
                 generator=gen_u,
-                device=X_u_w.device,
+                device=get_torch_device(X_u_w),
                 steps=steps_per_epoch,
             )
             for step, ((x_lb, y_lb), idx_u) in enumerate(zip(iter_l, iter_u_idx, strict=False)):
-                x_uw = X_u_w[idx_u]
-                x_us = X_u_s[idx_u]
+                x_uw = slice_data(X_u_w, idx_u)
+                x_us = slice_data(X_u_s, idx_u)
 
                 if bool(self.spec.use_cat):
-                    inputs = torch.cat([x_lb, x_uw, x_us], dim=0)
+                    inputs = concat_data([x_lb, x_uw, x_us])
                     logits = extract_logits(model(inputs))
                     if int(logits.ndim) != 2:
                         raise InductiveValidationError("Model logits must be 2D (batch, classes).")
-                    num_lb = int(x_lb.shape[0])
-                    num_u = int(x_uw.shape[0])
-                    expected = num_lb + num_u + int(x_us.shape[0])
+                    num_lb = int(get_torch_len(x_lb))
+                    num_u = int(get_torch_len(x_uw))
+                    expected = num_lb + num_u + int(get_torch_len(x_us))
                     if int(logits.shape[0]) != expected:
                         raise InductiveValidationError(
                             "Concatenated logits batch size does not match inputs."
@@ -350,28 +354,35 @@ class ADSHMethod(InductiveMethod):
         if backend != "torch":
             raise InductiveValidationError("ADSH predict_proba requires torch tensors.")
         torch = optional_import("torch", extra="inductive-torch")
-        if not isinstance(X, torch.Tensor):
-            raise InductiveValidationError("predict_proba requires torch.Tensor inputs.")
+        if not isinstance(X, torch.Tensor) and not isinstance(X, dict):
+            raise InductiveValidationError("predict_proba requires torch.Tensor or dict inputs.")
 
         model = self._bundle.model
         was_training = model.training
         model.eval()
 
         batch_size = int(self.spec.batch_size)
-        n_samples = int(X.shape[0])
+        n_samples = int(get_torch_len(X))
         all_logits = []
 
         with torch.no_grad():
             for start in range(0, n_samples, batch_size):
                 end = min(start + batch_size, n_samples)
-                batch_X = X[start:end]
+                if isinstance(X, dict):
+                    idx = torch.arange(start, end, device=get_torch_device(X))
+                    batch_X = slice_data(X, idx)
+                else:
+                    batch_X = X[start:end]
                 logits_batch = extract_logits(model(batch_X))
 
                 if int(logits_batch.ndim) != 2:
                     raise InductiveValidationError("Model logits must be 2D (batch, classes).")
                 all_logits.append(logits_batch)
 
-            logits = torch.cat(all_logits, dim=0)
+            if not all_logits:
+                logits = torch.empty((0, 0), device=get_torch_device(X))
+            else:
+                logits = torch.cat(all_logits, dim=0)
             proba = torch.softmax(logits, dim=1)
 
         if was_training:

{modssc-0.2.1.dev0 → modssc-0.2.2}/src/modssc/inductive/methods/co_training.py

@@ -102,6 +102,13 @@ def _get_torch_device(obj: Any) -> Any:
     return obj.device
 
 
+def _same_device(a: Any, b: Any) -> bool:
+    return (a == b) or (
+        getattr(a, "type", None) == getattr(b, "type", None)
+        and (getattr(a, "index", None) is None or getattr(b, "index", None) is None)
+    )
+
+
 def _view_payload_torch(value: Any, *, name: str):
     torch = optional_import("torch", extra="inductive-torch")
     if isinstance(value, Mapping):
@@ -280,7 +287,16 @@ class CoTrainingMethod(InductiveMethod):
         if backend == "torch":
             d1 = _get_torch_device(v1_l)
             d2 = _get_torch_device(v2_l)
-            if y_l.device != d1 or y_l.device != d2:
+            if d1 != d2:
+                raise InductiveValidationError("views must be on the same device.")
+            if not _same_device(y_l.device, d1):
+                try:
+                    y_l = y_l.to(d1)
+                except Exception as exc:
+                    raise InductiveValidationError(
+                        "y_l must be on the same device as the view tensors."
+                    ) from exc
+            if not _same_device(y_l.device, d1) or not _same_device(y_l.device, d2):
                 raise InductiveValidationError(
                     "y_l must be on the same device as the view tensors."
                 )

{modssc-0.2.1.dev0 → modssc-0.2.2}/src/modssc/inductive/methods/comatch.py

@@ -10,6 +10,7 @@ from modssc.inductive.base import InductiveMethod, MethodInfo
 from modssc.inductive.deep import TorchModelBundle
 from modssc.inductive.errors import InductiveValidationError
 from modssc.inductive.methods.deep_utils import (
+    concat_data,
     cycle_batch_indices,
     cycle_batches,
     ensure_float_tensor,
@@ -432,13 +433,13 @@ class CoMatchMethod(InductiveMethod):
                 x_us1 = _slice(X_u_s1, idx_u)
 
                 if bool(self.spec.use_cat):
-                    inputs = torch.cat([x_lb, x_uw, x_us0, x_us1], dim=0)
+                    inputs = concat_data([x_lb, x_uw, x_us0, x_us1])
                     logits_all, feats_all = _extract_logits_and_features(model(inputs))
                     if int(logits_all.ndim) != 2 or int(feats_all.ndim) != 2:
                         raise InductiveValidationError("Model logits/feat must be 2D tensors.")
-                    num_lb = int(x_lb.shape[0])
-                    num_u = int(x_uw.shape[0])
-                    expected = num_lb + num_u + int(x_us0.shape[0]) + int(x_us1.shape[0])
+                    num_lb = int(_len(x_lb))
+                    num_u = int(_len(x_uw))
+                    expected = num_lb + num_u + int(_len(x_us0)) + int(_len(x_us1))
                     if int(logits_all.shape[0]) != expected or int(feats_all.shape[0]) != expected:
                         raise InductiveValidationError(
                             "Concatenated logits/feat batch size does not match inputs."

{modssc-0.2.1.dev0 → modssc-0.2.2}/src/modssc/inductive/methods/daso.py

@@ -11,6 +11,7 @@ from modssc.inductive.base import InductiveMethod, MethodInfo
 from modssc.inductive.deep import TorchModelBundle
 from modssc.inductive.errors import InductiveValidationError
 from modssc.inductive.methods.deep_utils import (
+    concat_data,
     cycle_batch_indices,
     cycle_batches,
     ensure_float_tensor,
@@ -18,7 +19,10 @@ from modssc.inductive.methods.deep_utils import (
     ensure_model_device,
     extract_features,
     extract_logits,
+    get_torch_device,
+    get_torch_len,
     num_batches,
+    slice_data,
 )
 from modssc.inductive.methods.utils import (
     detect_backend,
@@ -257,15 +261,15 @@ class DASOMethod(InductiveMethod):
         X_u_s = ds.X_u_s
         logger.info(
             "DASO sizes: n_labeled=%s n_unlabeled=%s",
-            int(X_l.shape[0]),
-            int(X_u_w.shape[0]),
+            int(get_torch_len(X_l)),
+            int(get_torch_len(X_u_w)),
         )
 
-        if int(X_l.shape[0]) == 0:
+        if int(get_torch_len(X_l)) == 0:
             raise InductiveValidationError("X_l must be non-empty.")
-        if int(X_u_w.shape[0]) == 0 or int(X_u_s.shape[0]) == 0:
+        if int(get_torch_len(X_u_w)) == 0 or int(get_torch_len(X_u_s)) == 0:
             raise InductiveValidationError("X_u_w and X_u_s must be non-empty.")
-        if int(X_u_w.shape[0]) != int(X_u_s.shape[0]):
+        if int(get_torch_len(X_u_w)) != int(get_torch_len(X_u_s)):
             raise InductiveValidationError("X_u_w and X_u_s must have the same number of rows.")
 
         ensure_float_tensor(X_l, name="X_l")
@@ -280,7 +284,7 @@ class DASOMethod(InductiveMethod):
         bundle = ensure_model_bundle(self.spec.model_bundle)
         model = bundle.model
         optimizer = bundle.optimizer
-        ensure_model_device(model, device=X_l.device)
+        ensure_model_device(model, device=get_torch_device(X_l))
 
         if int(self.spec.batch_size) <= 0:
             raise InductiveValidationError("batch_size must be >= 1.")
@@ -317,8 +321,8 @@ class DASOMethod(InductiveMethod):
             ema_model.eval()
         self._ema_model = ema_model
 
-        steps_l = num_batches(int(X_l.shape[0]), int(self.spec.batch_size))
-        steps_u = num_batches(int(X_u_w.shape[0]), int(self.spec.batch_size))
+        steps_l = num_batches(int(get_torch_len(X_l)), int(self.spec.batch_size))
+        steps_u = num_batches(int(get_torch_len(X_u_w)), int(self.spec.batch_size))
         steps_per_epoch = max(int(steps_l), int(steps_u))
 
         gen_l = torch.Generator().manual_seed(int(seed))
@@ -339,24 +343,24 @@ class DASOMethod(InductiveMethod):
                 steps=steps_per_epoch,
             )
             iter_u_idx = cycle_batch_indices(
-                int(X_u_w.shape[0]),
+                int(get_torch_len(X_u_w)),
                 batch_size=int(self.spec.batch_size),
                 generator=gen_u,
-                device=X_u_w.device,
+                device=get_torch_device(X_u_w),
                 steps=steps_per_epoch,
             )
             for step, ((x_lb, y_lb), idx_u) in enumerate(zip(iter_l, iter_u_idx, strict=False)):
-                x_uw = X_u_w[idx_u]
-                x_us = X_u_s[idx_u]
+                x_uw = slice_data(X_u_w, idx_u)
+                x_us = slice_data(X_u_s, idx_u)
 
                 if bool(self.spec.use_cat):
-                    inputs = torch.cat([x_lb, x_uw, x_us], dim=0)
+                    inputs = concat_data([x_lb, x_uw, x_us])
                     logits_all, feats_all = _forward_logits_features(bundle, inputs)
                     if int(logits_all.ndim) != 2:
                         raise InductiveValidationError("Model logits must be 2D (batch, classes).")
-                    num_lb = int(x_lb.shape[0])
-                    num_u = int(x_uw.shape[0])
-                    expected = num_lb + num_u + int(x_us.shape[0])
+                    num_lb = int(get_torch_len(x_lb))
+                    num_u = int(get_torch_len(x_uw))
+                    expected = num_lb + num_u + int(get_torch_len(x_us))
                     if int(logits_all.shape[0]) != expected:
                         raise InductiveValidationError(
                             "Concatenated logits batch size does not match inputs."
@@ -375,8 +379,12 @@ class DASOMethod(InductiveMethod):
                     else:
                         logits_uw, feats_uw = _forward_logits_features(bundle, x_uw)
                     logits_us, feats_us = _forward_logits_features(bundle, x_us)
-                    feats_uw = _flatten_features(feats_uw, name="feat_uw", batch=int(x_uw.shape[0]))
-                    feats_us = _flatten_features(feats_us, name="feat_us", batch=int(x_us.shape[0]))
+                    feats_uw = _flatten_features(
+                        feats_uw, name="feat_uw", batch=int(get_torch_len(x_uw))
+                    )
+                    feats_us = _flatten_features(
+                        feats_us, name="feat_us", batch=int(get_torch_len(x_us))
+                    )
 
                 if int(logits_l.ndim) != 2 or int(logits_uw.ndim) != 2 or int(logits_us.ndim) != 2:
                     raise InductiveValidationError("Model logits must be 2D (batch, classes).")
@@ -409,7 +417,7 @@ class DASOMethod(InductiveMethod):
                 with torch.no_grad():
                     feat_lb = _forward_features(bundle, x_lb, model_override=ema_model)
                     feat_lb = _flatten_features(
-                        feat_lb, name="feat_lb", batch=int(x_lb.shape[0])
+                        feat_lb, name="feat_lb", batch=int(get_torch_len(x_lb))
                     ).detach()
                     if int(feat_lb.shape[1]) != feat_dim:
                         raise InductiveValidationError("Prototype feature dimension mismatch.")

{modssc-0.2.1.dev0 → modssc-0.2.2}/src/modssc/inductive/methods/deep_utils.py

@@ -82,6 +82,92 @@ def ensure_float_tensor(x: Any, *, name: str) -> None:
         raise InductiveValidationError(f"{name} must be float32 or float64.")
 
 
+def get_torch_len(x: Any) -> int:
+    if isinstance(x, dict) and "x" in x:
+        return int(x["x"].shape[0])
+    return int(x.shape[0])
+
+
+def get_torch_device(x: Any) -> Any:
+    if isinstance(x, dict) and "x" in x:
+        return x["x"].device
+    return x.device
+
+
+def get_torch_feature_dim(x: Any) -> int:
+    if isinstance(x, dict) and "x" in x:
+        return int(x["x"].shape[1])
+    return int(x.shape[1])
+
+
+def get_torch_ndim(x: Any) -> int:
+    if isinstance(x, dict) and "x" in x:
+        return int(x["x"].ndim)
+    return int(x.ndim)
+
+
+def concat_data(items: list[Any]) -> Any:
+    """Concatenate tensors or graph dicts along the batch dimension.
+
+    For dict inputs with 'x' and 'edge_index', this builds a disjoint union graph
+    by offsetting edge indices for each block.
+    """
+    if not items:
+        return items
+    if not any(isinstance(x, dict) for x in items):
+        torch = _torch()
+        return torch.cat(items, dim=0)
+
+    torch = _torch()
+    graphs = []
+    for x in items:
+        if isinstance(x, dict):
+            graphs.append(x)
+        else:
+            graphs.append({"x": x})
+
+    xs = [torch.as_tensor(g["x"]) for g in graphs]
+    x_cat = torch.cat(xs, dim=0)
+    out: dict[str, Any] = {"x": x_cat}
+
+    if all(isinstance(g, dict) and "edge_index" in g for g in graphs):
+        edge_indices = []
+        edge_weights = []
+        offset = 0
+        for g in graphs:
+            ei = g["edge_index"]
+            if not isinstance(ei, torch.Tensor):
+                ei = torch.as_tensor(ei, device=x_cat.device, dtype=torch.long)
+            ei = ei + offset
+            edge_indices.append(ei)
+            if "edge_weight" in g:
+                ew = g["edge_weight"]
+                if not isinstance(ew, torch.Tensor):
+                    ew = torch.as_tensor(ew, device=x_cat.device)
+                edge_weights.append(ew)
+            offset += int(g["x"].shape[0])
+        out["edge_index"] = torch.cat(edge_indices, dim=1)
+        if edge_weights:
+            out["edge_weight"] = torch.cat(edge_weights, dim=0)
+
+    # Concatenate per-node tensors for shared keys.
+    shared_keys = set.intersection(*(set(g.keys()) for g in graphs)) if graphs else set()
+    for key in shared_keys:
+        if key in {"x", "edge_index", "edge_weight"}:
+            continue
+        vals = [g[key] for g in graphs]
+        if all(isinstance(v, torch.Tensor) for v in vals):
+            # If node-wise (matching each graph's node count), concatenate
+            if all(v.shape[0] == g["x"].shape[0] for v, g in zip(vals, graphs, strict=False)):
+                out[key] = torch.cat(vals, dim=0)
+            else:
+                out[key] = vals[0]
+        else:
+            out[key] = vals[0]
+
+    return out
+
+
 @contextmanager
 def freeze_batchnorm(model: Any, *, enabled: bool):
     if not enabled:
@@ -142,8 +228,20 @@ def slice_data(X: Any, idx: Any) -> Any:
             try:
                 from torch_geometric.utils import subgraph
 
+                subgraph_idx = idx
+                if isinstance(idx, slice):
+                    start, stop, step = idx.indices(n)
+                    device = (
+                        X["edge_index"].device
+                        if isinstance(X["edge_index"], torch.Tensor)
+                        else None
+                    )
+                    subgraph_idx = torch.arange(start, stop, step, device=device)
+
                 # relabel_nodes=True ensures indices map to 0..len(idx)
-                edge_index, _ = subgraph(idx, X["edge_index"], relabel_nodes=True, num_nodes=n)
+                edge_index, _ = subgraph(
+                    subgraph_idx, X["edge_index"], relabel_nodes=True, num_nodes=n
+                )
                 batch_data["edge_index"] = edge_index
             except ImportError:
                 pass