Make typing imports optional with TYPE_CHECKING flag

pchlenski · pchlenski · commit 365d06bc4a37 · 2025-06-11T00:51:46.000-07:00
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -43,7 +43,7 @@ jobs:
         run: black --check manify/ --line-length 120
 
       - name: Check import ordering with isort
-        run: isort --check-only --profile black manify/ --line-width 120
+        run: isort --check-only --diff manify/
 
       - name: Run pylint
         run: pylint manify/
diff --git a/manify/clustering/fuzzy_kmeans.py b/manify/clustering/fuzzy_kmeans.py
@@ -20,14 +20,18 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import numpy as np
 import torch
-from beartype.typing import Literal
 from geoopt import ManifoldParameter
 from geoopt.optim import RiemannianAdam
-from jaxtyping import Float, Int
 from sklearn.base import BaseEstimator, ClusterMixin
 
+if TYPE_CHECKING:
+    from beartype.typing import Literal
+    from jaxtyping import Float, Int
+
 from ..manifolds import Manifold, ProductManifold
 from ..optimizers.radan import RiemannianAdan
 
diff --git a/manify/curvature_estimation/delta_hyperbolicity.py b/manify/curvature_estimation/delta_hyperbolicity.py
@@ -11,8 +11,12 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import torch
-from jaxtyping import Float, Int
+
+if TYPE_CHECKING:
+    from jaxtyping import Float, Int
 
 
 def sampled_delta_hyperbolicity(
diff --git a/manify/curvature_estimation/greedy_method.py b/manify/curvature_estimation/greedy_method.py
@@ -6,8 +6,12 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import torch
-from jaxtyping import Float
+
+if TYPE_CHECKING:
+    from jaxtyping import Float
 
 from ..manifolds import ProductManifold
 
diff --git a/manify/embedders/_base.py b/manify/embedders/_base.py
@@ -3,12 +3,15 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
 
 import torch
-from beartype.typing import Any
-from jaxtyping import Float
 from sklearn.base import BaseEstimator, TransformerMixin
 
+if TYPE_CHECKING:
+    from beartype.typing import Any
+    from jaxtyping import Float
+
 from ..manifolds import ProductManifold
 
 
diff --git a/manify/embedders/_losses.py b/manify/embedders/_losses.py
@@ -7,9 +7,13 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import networkx as nx
 import torch
-from jaxtyping import Float
+
+if TYPE_CHECKING:
+    from jaxtyping import Float
 
 from ..manifolds import ProductManifold
 
diff --git a/manify/embedders/coordinate_learning.py b/manify/embedders/coordinate_learning.py
@@ -11,12 +11,15 @@
 
 import sys
 import warnings
+from typing import TYPE_CHECKING
 
 import geoopt
 import numpy as np
 import torch
-from beartype.typing import Any
-from jaxtyping import Float, Int
+
+if TYPE_CHECKING:
+    from beartype.typing import Any
+    from jaxtyping import Float, Int
 
 from ..manifolds import ProductManifold
 from ._base import BaseEmbedder
diff --git a/manify/embedders/siamese.py b/manify/embedders/siamese.py
@@ -11,10 +11,13 @@
 from __future__ import annotations
 
 import sys
+from typing import TYPE_CHECKING
 
 import numpy as np
 import torch
-from jaxtyping import Float
+
+if TYPE_CHECKING:
+    from jaxtyping import Float
 
 from ..manifolds import ProductManifold
 from ._base import BaseEmbedder
diff --git a/manify/embedders/vae.py b/manify/embedders/vae.py
@@ -11,10 +11,13 @@
 from __future__ import annotations
 
 import sys
+from typing import TYPE_CHECKING
 
 import numpy as np
 import torch
-from jaxtyping import Float
+
+if TYPE_CHECKING:
+    from jaxtyping import Float
 
 from ..manifolds import ProductManifold
 from ._base import BaseEmbedder
diff --git a/manify/manifolds.py b/manify/manifolds.py
@@ -10,11 +10,14 @@
 from __future__ import annotations
 
 import warnings
+from typing import TYPE_CHECKING
 
 import geoopt
 import torch
-from beartype.typing import Callable, Literal
-from jaxtyping import Float, Real
+
+if TYPE_CHECKING:
+    from beartype.typing import Callable, Literal
+    from jaxtyping import Float, Real
 
 warnings.filterwarnings("ignore", category=UserWarning, module="torch.distributions")  # Singular samples from Wishart
 
diff --git a/manify/optimizers/_adan.py b/manify/optimizers/_adan.py
@@ -18,9 +18,12 @@
 """
 
 import math
-from beartype.typing import List
+from typing import TYPE_CHECKING
 
 import torch
+
+if TYPE_CHECKING:
+    from beartype.typing import List
 from torch import Tensor
 from torch.optim.optimizer import Optimizer
 
diff --git a/manify/optimizers/radan.py b/manify/optimizers/radan.py
@@ -36,10 +36,15 @@
 
 from __future__ import annotations
 
-from beartype.typing import Any, Callable
+from typing import TYPE_CHECKING
+
+import torch
 from geoopt import ManifoldParameter, ManifoldTensor
 from geoopt.optim.mixin import OptimMixin
-from jaxtyping import Float
+
+if TYPE_CHECKING:
+    from beartype.typing import Any, Callable
+    from jaxtyping import Float
 
 from . import _adan
 
diff --git a/manify/predictors/_base.py b/manify/predictors/_base.py
@@ -3,12 +3,15 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
 
 import torch
-from beartype.typing import Literal
-from jaxtyping import Float
 from sklearn.base import BaseEstimator
 
+if TYPE_CHECKING:
+    from beartype.typing import Literal
+    from jaxtyping import Float
+
 from ..manifolds import ProductManifold
 
 
diff --git a/manify/predictors/_kernel.py b/manify/predictors/_kernel.py
@@ -2,8 +2,12 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import torch
-from jaxtyping import Float
+
+if TYPE_CHECKING:
+    from jaxtyping import Float
 
 from ..manifolds import Manifold, ProductManifold
 
diff --git a/manify/predictors/_midpoint.py b/manify/predictors/_midpoint.py
@@ -2,8 +2,12 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import torch
-from jaxtyping import Float
+
+if TYPE_CHECKING:
+    from jaxtyping import Float
 
 from ..manifolds import Manifold
 
diff --git a/manify/predictors/decision_tree.py b/manify/predictors/decision_tree.py
@@ -5,11 +5,15 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING, List, Optional, Tuple
+
 import torch
-from beartype.typing import Any, Literal
-from jaxtyping import Bool, Float, Int, Real
 from sklearn.base import BaseEstimator, ClassifierMixin
 
+if TYPE_CHECKING:
+    from beartype.typing import Any, Literal
+    from jaxtyping import Bool, Float, Int, Real
+
 from ..manifolds import ProductManifold
 from ._midpoint import midpoint
 
diff --git a/manify/predictors/kappa_gcn.py b/manify/predictors/kappa_gcn.py
@@ -3,11 +3,14 @@
 from __future__ import annotations
 
 import sys
+from typing import TYPE_CHECKING, Optional
 
 import geoopt
 import torch
-from beartype.typing import Callable, Literal
-from jaxtyping import Float, Real
+
+if TYPE_CHECKING:
+    from beartype.typing import Callable, Literal
+    from jaxtyping import Float, Real
 
 from ..manifolds import Manifold, ProductManifold
 
diff --git a/manify/predictors/perceptron.py b/manify/predictors/perceptron.py
@@ -2,10 +2,14 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import torch
-from jaxtyping import Float, Int
 from sklearn.base import BaseEstimator, ClassifierMixin
 
+if TYPE_CHECKING:
+    from jaxtyping import Float, Int
+
 from ..manifolds import ProductManifold
 from ._kernel import product_kernel
 
diff --git a/manify/predictors/svm.py b/manify/predictors/svm.py
@@ -2,13 +2,17 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import cvxpy
 import numpy as np
 import torch
-from beartype.typing import Literal
-from jaxtyping import Float, Int
 from sklearn.base import BaseEstimator, ClassifierMixin
 
+if TYPE_CHECKING:
+    from beartype.typing import Literal
+    from jaxtyping import Float, Int
+
 from ..manifolds import ProductManifold
 from ._kernel import product_kernel
 
diff --git a/manify/utils/benchmarks.py b/manify/utils/benchmarks.py
@@ -3,12 +3,15 @@
 from __future__ import annotations
 
 import time
+from typing import TYPE_CHECKING, Dict
 
 import numpy as np
 import torch
-from beartype.typing import Literal, TypeAlias
-from jaxtyping import Float, Real
 from sklearn.base import BaseEstimator
+
+if TYPE_CHECKING:
+    from beartype.typing import Literal, TypeAlias
+    from jaxtyping import Float, Real
 from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
 from sklearn.linear_model import SGDClassifier, SGDRegressor
 from sklearn.metrics import accuracy_score, f1_score, mean_squared_error, root_mean_squared_error
diff --git a/manify/utils/dataloaders.py b/manify/utils/dataloaders.py
@@ -33,9 +33,13 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING, Optional
+
 import torch
 from datasets import load_dataset
-from jaxtyping import Float, Real
+
+if TYPE_CHECKING:
+    from jaxtyping import Float, Real
 
 
 def load_hf(name: str, namespace: str = "manify") -> tuple[
diff --git a/manify/utils/link_prediction.py b/manify/utils/link_prediction.py
@@ -2,11 +2,15 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import torch
-from beartype.typing import Any
-from jaxtyping import Float, Int
 from sklearn.model_selection import train_test_split
 
+if TYPE_CHECKING:
+    from beartype.typing import Any
+    from jaxtyping import Float, Int
+
 from ..manifolds import ProductManifold
 
 
diff --git a/manify/utils/preprocessing.py b/manify/utils/preprocessing.py
@@ -2,8 +2,12 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import torch
-from jaxtyping import Float
+
+if TYPE_CHECKING:
+    from jaxtyping import Float
 
 
 def knn_graph(x: Float[torch.Tensor, "n_points n_dim"], k: int) -> Float[torch.Tensor, "n_points n_points"]:
diff --git a/manify/utils/visualization.py b/manify/utils/visualization.py
@@ -2,8 +2,12 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import torch
-from jaxtyping import Float
+
+if TYPE_CHECKING:
+    from jaxtyping import Float
 
 
 def hyperboloid_to_poincare(X: Float[torch.Tensor, "n_points n_dim"]) -> Float[torch.Tensor, "n_points n_dim-1"]:
diff --git a/notebooks/65_fix_time_series.ipynb b/notebooks/65_fix_time_series.ipynb
@@ -1142,7 +1142,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.0"
+   "version": "undefined.undefined.undefined"
   }
  },
  "nbformat": 4,
