Fix jaxtyping compatibiilty; isort benchmarks

Author: pchlenski

manify/__init__.py

@@ -3,6 +3,9 @@
 import manify.manifolds
 import manify.predictors
 
+# Define version
+__version__ = "0.0.2"
+
 # Dynamically check for utils dependencies
 try:
     import importlib.util

manify/manifolds.py

@@ -11,7 +11,8 @@
 from __future__ import annotations
 
 import warnings
-from typing import TYPE_CHECKING, Callable, List, Literal, Optional, Tuple, Union
+from typing import (TYPE_CHECKING, Callable, List, Literal, Optional, Tuple,
+                    Union)
 
 import geoopt
 import torch

manify/utils/benchmarks.py

@@ -9,24 +9,19 @@
 from sklearn.base import BaseEstimator
 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,
-)
+from sklearn.metrics import (accuracy_score, f1_score, mean_squared_error,
+                             root_mean_squared_error)
 from sklearn.model_selection import train_test_split
 from sklearn.neighbors import KNeighborsClassifier, KNeighborsRegressor
 from sklearn.svm import SVC, SVR
 from sklearn.tree import DecisionTreeClassifier, DecisionTreeRegressor
 
 from ..manifolds import ProductManifold
+from ..predictors.decision_tree import ProductSpaceDT, ProductSpaceRF
 from ..predictors.kappa_gcn import KappaGCN, get_A_hat
 from ..predictors.perceptron import ProductSpacePerceptron
 from ..predictors.svm import ProductSpaceSVM
 
-from ..predictors.decision_tree import ProductSpaceDT, ProductSpaceRF
-
 
 def _score(
     _X: Float[torch.Tensor, "n_samples n_dims"],

pyproject.toml

@@ -23,6 +23,7 @@ dependencies = [
 [project.optional-dependencies]
 dev = [
     "jaxtyping",
+    "beartype",
     "pytest",
     "mypy",
     "pytest-cov",
@@ -70,4 +71,14 @@ generated-members = "torch.*,nn.*"
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 addopts = "--jaxtyping-packages=beartype.beartype"
-python_files = "test_*.py"
+python_files = "test_*.py"
+
+[tool.mypy]
+# Specify the packages to check
+packages = ["manify"]
+
+# Ignore missing imports for external libraries
+ignore_missing_imports = true
+
+# Use jaxtyping-specific plugins
+plugins = ["jaxtyping.mypy_plugin"]