Expand distributed indexing, match numpy indexing scheme

heat/classification/kneighborsclassifier.py

@@ -122,11 +122,11 @@ def predict(self, x: DNDarray) -> DNDarray:
         """
         distances = self.effective_metric_(x, self.x)
         _, indices = ht.topk(distances, self.n_neighbors, largest=False)
-        predictions = self.y[indices.flatten()]
+
+        predictions = self.y[indices]
         predictions.balance_()
-        predictions = ht.reshape(predictions, (indices.gshape + (self.y.gshape[1],)))
+        predictions = ht.reshape(predictions, indices.gshape + (self.y.gshape[1],))
         predictions = ht.sum(predictions, axis=1)
 
         self.classes_ = ht.argmax(predictions, axis=1)
-
         return self.classes_

heat/cluster/batchparallelclustering.py

@@ -42,7 +42,25 @@ def _initialize_plus_plus(
     for i in range(1, n_clusters):
         dist = torch.cdist(X, X[idxs[:i]], p=p)
         dist = torch.min(dist, dim=1)[0]
-        idxs[i] = torch.multinomial(weights * dist, 1)
+        probs = weights * dist
+        probs = torch.nan_to_num(probs, nan=0.0, posinf=0.0, neginf=0.0)
+
+        # Minimal fallback ONLY if multinomial would crash
+        if probs.sum() <= 0:
+            # fall back to standard k-means++ (ignore weights)
+            probs = torch.nan_to_num(dist, nan=0.0, posinf=0.0, neginf=0.0)
+
+        if probs.sum() <= 0:
+            # fully degenerate (all distances zero) -> pick any not-yet-picked index if possible
+            mask = torch.ones(X.shape[0], dtype=torch.bool, device=X.device)
+            mask[idxs[:i]] = False
+            candidates = torch.nonzero(mask, as_tuple=False).flatten()
+            if candidates.numel() > 0:
+                idxs[i] = candidates[torch.randint(0, candidates.numel(), (1,), device=X.device)]
+            else:
+                idxs[i] = torch.randint(0, X.shape[0], (1,), device=X.device)
+        else:
+            idxs[i] = torch.multinomial(probs, 1)
     return X[idxs]