Fix jaxtyping annotations

Author: pchlenski

manify/__init__.py

@@ -1,5 +1,9 @@
 """Manify: A Python Library for Learning Non-Euclidean Representations."""
 
+from jaxtyping import install_import_hook
+
+install_import_hook("manify", "beartype.beartype")
+
 from manify.curvature_estimation import (
     delta_hyperbolicity,
     greedy_signature_selection,

manify/curvature_estimation/__init__.py

@@ -7,8 +7,13 @@
 * `sectional_curvature`: Estimates the sectional curvature of a graph from its distance matrix.
 """
 
-from manify.curvature_estimation.delta_hyperbolicity import delta_hyperbolicity, sampled_delta_hyperbolicity
+from manify.curvature_estimation.delta_hyperbolicity import vectorized_delta_hyperbolicity, sampled_delta_hyperbolicity
 from manify.curvature_estimation.greedy_method import greedy_signature_selection
 from manify.curvature_estimation.sectional_curvature import sectional_curvature
 
-__all__ = ["greedy_signature_selection", "sectional_curvature", "sampled_delta_hyperbolicity", "delta_hyperbolicity"]
+__all__ = [
+    "greedy_signature_selection",
+    "sectional_curvature",
+    "sampled_delta_hyperbolicity",
+    "vectorized_delta_hyperbolicity",
+]

manify/curvature_estimation/delta_hyperbolicity.py

@@ -19,7 +19,7 @@
 
 def sampled_delta_hyperbolicity(
     D: Float[torch.Tensor, "n_points n_points"], n_samples: int = 1000, reference_idx: int = 0, relative: bool = True
-) -> Tuple[Float[torch.Tensor, "n_samples,"], Float[torch.Tensor, "n_samples 3"]]:
+) -> Tuple[Float[torch.Tensor, "n_samples"], Float[torch.Tensor, "n_samples 3"]]:
     r"""Computes $\delta$-hyperbolicity by sampling random point triplets.
 
     For large metric spaces, this approximates $\delta$-hyperbolicity by randomly sampling triplets. For each triplet
@@ -61,7 +61,7 @@ def sampled_delta_hyperbolicity(
     return deltas, indices
 
 
-def delta_hyperbolicity(
+def vectorized_delta_hyperbolicity(
     D: Float[torch.Tensor, "n_points n_points"], reference_idx: int = 0, relative: bool = True, full: bool = False
 ) -> Float[torch.Tensor, "n_points n_points n_points"]:
     r"""Computes the exact delta-hyperbolicity of a metric space over all point triplets.

manify/curvature_estimation/greedy_method.py

@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 from typing import Any, Tuple
+from jaxtyping import Float
 
 import torch
 
@@ -15,7 +16,7 @@
 
 def greedy_signature_selection(
     pm: ProductManifold,
-    dists: torch.Tensor,
+    dists: Float[torch.Tensor, "n_points n_points"],
     candidate_components: Tuple[Tuple[float, int], ...] = ((-1.0, 2), (0.0, 2), (1.0, 2)),
     max_components: int = 3,
 ) -> Any:

manify/embedders/_base.py

@@ -13,7 +13,7 @@
 
 
 class BaseEmbedder(BaseEstimator, TransformerMixin, ABC):
-    """Base class for everything in manify.embedders.
+    """Base class for everything in `manify.embedders`.
 
     This is an abstract class that that defines a common interface for all embedding methods. We assume only that a
     ProductManifold object is given. We try to follow the scikit-learn API's fit/transform paradigm as closely as

manify/embedders/coordinate_learning.py

@@ -70,7 +70,7 @@ def fit(  # type: ignore[override]
         self,
         X: None,
         D: 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([]),
         lr: float = 1e-2,
         burn_in_lr: float = 1e-3,
         curvature_lr: float = 0.0,  # Off by default

manify/embedders/siamese.py

@@ -118,7 +118,7 @@ def forward(
     ) -> Tuple[
         Float[torch.Tensor, "batch_size n_latent"],
         Float[torch.Tensor, "batch_size n_latent"],
-        Float[torch.Tensor, "batch_size,"],
+        Float[torch.Tensor, "batch_size"],
         Float[torch.Tensor, "batch_size n_features"],
         Float[torch.Tensor, "batch_size n_features"],
     ]:

manify/embedders/vae.py

@@ -182,7 +182,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"]:
         r"""Computes the KL divergence between posterior and prior distributions in the manifold.
 
         For distributions in Riemannian manifolds, computing the KL divergence analytically
@@ -358,8 +358,8 @@ def fit(  # type: ignore[override]
         return self
 
     def transform(
-        self, X: Float["n_points n_features"], D: None = None, batch_size: int = 32, expmap: bool = True
-    ) -> Float["n_points embedding_dim"]:
+        self, X: Float[torch.Tensor, "n_points n_features"], D: None = None, batch_size: int = 32, expmap: bool = True
+    ) -> Float[torch.Tensor, "n_points embedding_dim"]:
         """Transform data using the trained VAE. Outputs means of the variational distribution.
 
         Args:

manify/manifolds.py

@@ -262,7 +262,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"]:
         r"""Compute probability density function for $\mathcal{WN}(\mathbf{z}; \mu, \Sigma)$ on the manifold.
 
         Args:
@@ -646,9 +646,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"]:
         r"""Compute probability density function for $\mathcal{WN}(\mathbf{z} ; \mu, \Sigma)$ on the product manifold.
 
         Args:
@@ -756,7 +756,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.
 
         Args:

manify/predictors/_base.py

@@ -0,0 +1,87 @@
+"""Base predictor class."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional
+
+import torch
+from jaxtyping import Float
+from sklearn.base import BaseEstimator
+
+from ..manifolds import ProductManifold
+
+
+class BasePredictor(BaseEstimator, ABC):
+    """Base class for everything in `manify.predictors`.
+
+    This is an abstract class that defines a common interface for all mixed-curvature predictors. We assume only that a
+    ProductManifold object is given. We try to follow the scikit-learn API's fit/predict_proba/predict paradigm as
+    closely as possible, while accommodating the nuances of product manifold geometry and Pytorch/Geoopt.
+
+    Attributes:
+        pm: ProductManifold object associated with the predictor.
+        task: Task type, either "classification" or "regression".
+        random_state: Random state for reproducibility.
+        device: Device for tensor computations. If not provided, defaults to pm.device.
+        loss_history_: History of loss values during training.
+        is_fitted_: Boolean flag indicating if the predictor has been fitted.
+    """
+
+    def __init__(
+        self,
+        pm: ProductManifold,
+        task: Literal["classification", "regression"],
+        random_state: Optional[int] = None,
+        device: Optional[str] = None,
+    ) -> None:
+        self.pm = pm
+        self.task = task
+        self.random_state = random_state
+        self.device = pm.device if device is None else device
+        self.loss_history_: Dict[str, List[float]] = {}
+        self.is_fitted_: bool = False
+
+    @abstractmethod
+    def fit(
+        self, X: Float[torch.Tensor, "n_points n_features"], y: Float[torch.Tensor, "n_points n_classes"]
+    ) -> "BasePredictor":
+        """Abstract method to fit a predictor. Requires features and labels.
+
+        Args:
+            X: Features to fit.
+            y: Labels for the features.
+
+        Returns:
+            self: Fitted predictor instance.
+        """
+        pass
+
+    @abstractmethod
+    def predict_proba(
+        self, X: Optional[Float[torch.Tensor, "n_points n_features"]]
+    ) -> Float[torch.Tensor, "n_points n_classes"]:
+        """Compute the predicted probabilities for the given features.
+
+        Args:
+            X: New inputs for which to make predictions.
+
+        Returns:
+            X_proba: Predicted probabilities for the input features.
+        """
+        pass
+
+    def predict(
+        self, X: Optional[Float[torch.Tensor, "n_points n_features"]]
+    ) -> Float[torch.Tensor, "n_points n_classes"]:
+        """Compute the predicted classes for the given features.
+
+        Args:
+            X: New inputs for which to make predictions.
+
+        Returns:
+            X_proba: Predicted probabilities for the input features.
+        """
+        if self.task == "regression":
+            return self.predict_proba(X=X)
+        return self.predict_proba(X=X).argmax(dim=-1)