Fix types everywere

Author: pchlenski

.github/workflows/test.yml

@@ -12,7 +12,7 @@ jobs:
     strategy:
       fail-fast: false    # don’t stop the matrix if one Python version fails
       matrix:
-        python-version: ["3.9", "3.10", "3.11"]
+        python-version: ["3.10", "3.11"] # jaxtyping requires >= 3.10; scipy requires < 3.12
 
     steps:
       # Setup and installation

manify/curvature_estimation/delta_hyperbolicity.py

@@ -29,7 +29,9 @@ def sampled_delta_hyperbolicity(dismat, n_samples=1000, reference_idx=0):
     return rel_deltas, indices
 
 
-def iterative_delta_hyperbolicity(dismat):
+def iterative_delta_hyperbolicity(
+    dismat: Float[torch.Tensor, "n_points n_points"],
+) -> Float[torch.Tensor, "n_points n_points n_points"]:
     """delta(x,y,z) = min((x,y)_w,(y-z)_w) - (x,z)_w"""
     n = dismat.shape[0]
     w = 0
@@ -56,7 +58,7 @@ def iterative_delta_hyperbolicity(dismat):
     return rel_deltas, gromov_products
 
 
-def gromov_product(i, j, k, dismat):
+def gromov_product(i: int, j: int, k: int, dismat: Float[torch.Tensor, "n_points n_points"]) -> float:
     """(j,k)_i = 0.5 (d(i,j) + d(i,k) - d(j,k))"""
     d_ij = dismat[i, j]
     d_ik = dismat[i, k]
@@ -65,18 +67,14 @@ def gromov_product(i, j, k, dismat):
 
 
 def delta_hyperbolicity(
-    dismat: Float[torch.Tensor, "n_points n_points"],
-    relative=True,
-    device="cpu",
-    full=False,
+    dismat: Float[torch.Tensor, "n_points n_points"], relative=True, full=False
 ) -> Float[torch.Tensor, ""]:
     """
     Compute the delta-hyperbolicity of a metric space.
 
     Args:
         dismat: Distance matrix of the metric space.
         relative: Whether to return the relative delta-hyperbolicity.
-        device: Device to run the computation on.
         full: Whether to return the full delta tensor or just the maximum delta.
 
     Returns:

manify/curvature_estimation/sectional_curvature.py

@@ -61,7 +61,7 @@ def sample(D, size, n_samples=100):
 def estimate(D, size, n_samples):
     samples = sample(D, size, n_samples)
     m1 = np.mean(samples)
-    m2 = np.mean(samples ** 2)
+    m2 = np.mean(samples**2)
     return samples
 
 

manify/embedders/coordinate_learning.py

@@ -23,7 +23,7 @@
 def train_coords(
     pm: ProductManifold,
     dists: Float[torch.Tensor, "n_points n_points"],
-    test_indices: Int[torch.Tensor, "n_test"] = torch.tensor([]),
+    test_indices: Int[torch.Tensor, "n_test,"] = torch.tensor([]),
     device: str = "cpu",
     burn_in_learning_rate: float = 1e-3,
     burn_in_iterations: int = 2_000,
@@ -32,7 +32,6 @@ def train_coords(
     training_iterations: int = 18_000,
     loss_window_size: int = 100,
     logging_interval: int = 10,
-    scale=1.0,
 ) -> Tuple[Float[torch.Tensor, "n_points n_dim"], Dict[str, List[float]]]:
     """
     Coordinate training and optimization
@@ -79,7 +78,7 @@ def train_coords(
     my_tqdm = tqdm(total=burn_in_iterations + training_iterations, leave=False)
 
     # Outer training loop - mostly setting optimizer learning rates up here
-    losses = {"train_train": [], "test_test": [], "train_test": [], "total": []}
+    losses: Dict[str, List[float]] = {"train_train": [], "test_test": [], "train_test": [], "total": []}
 
     # Actual training loop
     for i in range(burn_in_iterations + training_iterations):

manify/embedders/siamese.py

@@ -25,7 +25,7 @@ def __init__(
         if decoder is not None:
             self.decoder = decoder
         else:
-            self.decoder = lambda x: x
+            self.decoder = torch.nn.Identity()
             self.decoder.requires_grad_(False)
             self.decoder.to(pm.device)
 

manify/embedders/vae.py

@@ -40,9 +40,10 @@ def __init__(
         else:
             raise ValueError(f"Unknown reconstruction loss: {reconstruction_loss}")
 
-    def encode(
-        self, x: Float[torch.Tensor, "batch_size n_features"]
-    ) -> Tuple[Float[torch.Tensor, "batch_size n_latent"], Float[torch.Tensor, "batch_size n_latent"],]:
+    def encode(self, x: Float[torch.Tensor, "batch_size n_features"]) -> Tuple[
+        Float[torch.Tensor, "batch_size n_latent"],
+        Float[torch.Tensor, "batch_size n_latent"],
+    ]:
         """Must return z_mean, z_logvar"""
         z_mean_tangent, z_logvar = self.encoder(x)
         z_mean_ambient = z_mean_tangent @ self.pm.projection_matrix  # Adds zeros in the right places
@@ -53,9 +54,7 @@ def decode(self, z: Float[torch.Tensor, "batch_size n_latent"]) -> Float[torch.T
         """Decoding in product space VAE"""
         return self.decoder(z)
 
-    def forward(
-        self, x: Float[torch.Tensor, "batch_size n_features"]
-    ) -> Tuple[
+    def forward(self, x: Float[torch.Tensor, "batch_size n_features"]) -> Tuple[
         Float[torch.Tensor, "batch_size n_features"],
         Float[torch.Tensor, "batch_size n_latent"],
         List[Float[torch.Tensor, "n_latent n_latent"]],
@@ -82,7 +81,7 @@ def kl_divergence(
         self,
         z_mean: Float[torch.Tensor, "batch_size n_latent"],
         sigma_factorized: List[Float[torch.Tensor, "n_latent n_latent"]],
-    ) -> Float[torch.Tensor, "batch_size"]:
+    ) -> Float[torch.Tensor, "batch_size,"]:
         """
         Computes the KL divergence between posterior and prior distributions.
 

manify/manifolds.py

@@ -11,7 +11,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import TYPE_CHECKING, Callable, List, Literal, Optional, Tuple, Union
+from typing import Callable, List, Literal, Optional, Tuple, Union
 
 import geoopt
 import torch
@@ -33,13 +33,7 @@ class Manifold:
     stereographic: (bool) Whether to use stereographic coordinates for the manifold.
     """
 
-    def __init__(
-        self,
-        curvature: float,
-        dim: int,
-        device: str = "cpu",
-        stereographic: bool = False,
-    ):
+    def __init__(self, curvature: float, dim: int, device: str = "cpu", stereographic: bool = False):
         # Device management
         self.device = device
 
@@ -103,9 +97,7 @@ def to(self, device: str) -> "Manifold":
         return self
 
     def inner(
-        self,
-        X: Float[torch.Tensor, "n_points1 n_dim"],
-        Y: Float[torch.Tensor, "n_points2 n_dim"],
+        self, X: Float[torch.Tensor, "n_points1 n_dim"], Y: Float[torch.Tensor, "n_points2 n_dim"]
     ) -> Float[torch.Tensor, "n_points1 n_points2"]:
         """
         Compute the inner product of manifolds.
@@ -126,9 +118,7 @@ def inner(
         return X_fixed @ Y.T * scaler
 
     def dist(
-        self,
-        X: Float[torch.Tensor, "n_points1 n_dim"],
-        Y: Float[torch.Tensor, "n_points2 n_dim"],
+        self, X: Float[torch.Tensor, "n_points1 n_dim"], Y: Float[torch.Tensor, "n_points2 n_dim"]
     ) -> Float[torch.Tensor, "n_points1 n_points2"]:
         """
         Inherit distance function from the geoopt manifold.
@@ -143,9 +133,7 @@ def dist(
         return self.manifold.dist(X[:, None], Y[None, :])
 
     def dist2(
-        self,
-        X: Float[torch.Tensor, "n_points1 n_dim"],
-        Y: Float[torch.Tensor, "n_points2 n_dim"],
+        self, X: Float[torch.Tensor, "n_points1 n_dim"], Y: Float[torch.Tensor, "n_points2 n_dim"]
     ) -> Float[torch.Tensor, "n_points1 n_points2"]:
         """
         Inherit squared distance function from the geoopt manifold.
@@ -265,7 +253,7 @@ def log_likelihood(
         z: Float[torch.Tensor, "n_points n_ambient_dim"],
         mu: Optional[Float[torch.Tensor, "n_points n_ambient_dim"]] = None,
         sigma: Optional[Float[torch.Tensor, "n_points n_dim n_dim"]] = None,
-    ) -> Float[torch.Tensor, "n_points"]:
+    ) -> Float[torch.Tensor, "n_points,"]:
         """
         Probability density function for WN(z ; mu, Sigma) in manifold
 
@@ -321,9 +309,7 @@ def log_likelihood(
             return ll - (n - 1) * torch.log(R * torch.abs(sin_M(u_norm / R) / u_norm) + 1e-8)
 
     def logmap(
-        self,
-        x: Float[torch.Tensor, "n_points n_dim"],
-        base: Optional[Float[torch.Tensor, "n_points n_dim"]] = None,
+        self, x: Float[torch.Tensor, "n_points n_dim"], base: Optional[Float[torch.Tensor, "n_points n_dim"]] = None
     ) -> Float[torch.Tensor, "n_points n_dim"]:
         """
         Logarithmic map of point on manifold x at base point.
@@ -341,9 +327,7 @@ def logmap(
         return self.manifold.logmap(x=base, y=x)
 
     def expmap(
-        self,
-        u: Float[torch.Tensor, "n_points n_dim"],
-        base: Optional[Float[torch.Tensor, "n_points n_dim"]] = None,
+        self, u: Float[torch.Tensor, "n_points n_dim"], base: Optional[Float[torch.Tensor, "n_points n_dim"]] = None
     ) -> Float[torch.Tensor, "n_points n_dim"]:
         """
         Exponential map of tangent vector u at base point.
@@ -415,7 +399,7 @@ def inverse_stereographic(self, *points: Float[torch.Tensor, "n_points n_dim_ste
             return orig_manifold, *points  # type: ignore
 
         # Inverse projection for points
-        norm_squared = [(Y ** 2).sum(dim=1, keepdim=True) for Y in points]
+        norm_squared = [(Y**2).sum(dim=1, keepdim=True) for Y in points]
         sign = torch.sign(self.curvature)  # type: ignore
 
         X0 = (1 + sign * norm_squared) / (1 - sign * norm_squared)
@@ -457,12 +441,7 @@ class ProductManifold(Manifold):
     stereographic: (bool) Whether to use stereographic coordinates for the manifold.
     """
 
-    def __init__(
-        self,
-        signature: List[Tuple[float, int]],
-        device: str = "cpu",
-        stereographic: bool = False,
-    ):
+    def __init__(self, signature: List[Tuple[float, int]], device: str = "cpu", stereographic: bool = False):
         # Device management
         self.device = device
 
@@ -485,12 +464,7 @@ def __init__(
 
         # Manifold <-> Dimension mapping
         self.ambient_dim, self.n_manifolds, self.dim = 0, 0, 0
-        self.dim2man, self.man2dim, self.man2intrinsic, self.intrinsic2man = (
-            {},
-            {},
-            {},
-            {},
-        )
+        self.dim2man, self.man2dim, self.man2intrinsic, self.intrinsic2man = {}, {}, {}, {}
 
         for M in self.P:
             for d in range(self.ambient_dim, self.ambient_dim + M.ambient_dim):
@@ -549,10 +523,7 @@ def sample(
         sigma_factorized: Optional[List[Float[torch.Tensor, "n_points n_dim_manifold n_dim_manifold"]]] = None,
     ) -> Union[
         Float[torch.Tensor, "n_points n_ambient_dim"],
-        Tuple[
-            Float[torch.Tensor, "n_points n_ambient_dim"],
-            Float[torch.Tensor, "n_points n_dim"],
-        ],
+        Tuple[Float[torch.Tensor, "n_points n_ambient_dim"], Float[torch.Tensor, "n_points n_dim"]],
     ]:
         """
         Sample from the variational distribution.
@@ -593,9 +564,9 @@ def sample(
     def log_likelihood(
         self,
         z: Float[torch.Tensor, "batch_size n_dim"],
-        mu: Optional[Float[torch.Tensor, "n_dim"]] = None,
+        mu: Optional[Float[torch.Tensor, "n_dim,"]] = None,
         sigma_factorized: Optional[List[Float[torch.Tensor, "n_points n_dim_manifold n_dim_manifold"]]] = None,
-    ) -> Float[torch.Tensor, "batch_size"]:
+    ) -> Float[torch.Tensor, "batch_size,"]:
         """
         Probability density function for WN(z ; mu, Sigma) in manifold
 
@@ -624,7 +595,7 @@ def log_likelihood(
         ]
         return torch.cat(component_lls, axis=1).sum(axis=1)  # type: ignore
 
-    def stereographic(self, *points: Float[torch.Tensor, "n_points n_dim"]) -> Tuple[Manifold, ...]:
+    def stereographic(self, *points: Float[torch.Tensor, "n_points n_dim"]) -> Tuple["ProductManifold", ...]:
         if self.is_stereographic:
             print("Manifold is already in stereographic coordinates.")
             return self, *points  # type: ignore
@@ -653,7 +624,7 @@ def gaussian_mixture(
         regression_noise_std: float = 0.1,
         task: Literal["classification", "regression"] = "classification",
         adjust_for_dims: bool = False,
-    ) -> Tuple[Float[torch.Tensor, "n_points n_ambient_dim"], Float[torch.Tensor, "n_points"]]:
+    ) -> Tuple[Float[torch.Tensor, "n_points n_ambient_dim"], Float[torch.Tensor, "n_points,"]]:
         """
         Generate a set of labeled samples from a Gaussian mixture model.