nystrom-ncut 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

nystrom_ncut/__init__.py

@@ -1,4 +1,7 @@
-from .ncut_pytorch import NCUT
+from .ncut_pytorch import (
+    NCUT,
+    axis_align,
+)
 from .propagation_utils import (
     affinity_from_features,
     propagate_eigenvectors,
@@ -6,7 +9,6 @@ from .propagation_utils import (
     quantile_normalize,
 )
 from .visualize_utils import (
-    eigenvector_to_rgb,
     rgb_from_tsne_3d,
     rgb_from_umap_sphere,
     rgb_from_tsne_2d,
@@ -18,5 +20,3 @@ from .visualize_utils import (
     propagate_rgb_color,
     get_mask,
 )
-from .ncut_pytorch import nystrom_ncut, ncut
-from .ncut_pytorch import kway_ncut, axis_align

nystrom_ncut/common.py

@@ -0,0 +1,20 @@
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn.functional as Fn
+
+
+def ceildiv(a: int, b: int) -> int:
+    return -(-a // b)
+
+
+def lazy_normalize(x: torch.Tensor, n: int = 1000, **normalize_kwargs: Any) -> torch.Tensor:
+    numel = np.prod(x.shape[:-1])
+    n = min(n, numel)
+    random_indices = torch.randperm(numel)[:n]
+    _x = x.flatten(0, -2)[random_indices]
+    if torch.allclose(torch.norm(_x, **normalize_kwargs), torch.ones(n, device=x.device)):
+        return x
+    else:
+        return Fn.normalize(x, **normalize_kwargs)