Update to geoopt 0.5.1; add support for Python 3.12+

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.10", "3.11"] # jaxtyping requires >= 3.10; scipy requires < 3.12
+        python-version: ["3.10", "3.11", "3.12", "3.13"] # jaxtyping requires >= 3.10
 
     steps:
       # Setup and installation

manify/optimizers/_adan.py

@@ -22,8 +22,6 @@
 
 import torch
 
-if TYPE_CHECKING:
-    from beartype.typing import List
 from torch import Tensor
 from torch.optim.optimizer import Optimizer
 
@@ -242,12 +240,12 @@ def step(self, closure=None):
 
 
 def _single_tensor_adan(
-    params: List[Tensor],
-    grads: List[Tensor],
-    exp_avgs: List[Tensor],
-    exp_avg_sqs: List[Tensor],
-    exp_avg_diffs: List[Tensor],
-    neg_pre_grads: List[Tensor],
+    params: list[Tensor],
+    grads: list[Tensor],
+    exp_avgs: list[Tensor],
+    exp_avg_sqs: list[Tensor],
+    exp_avg_diffs: list[Tensor],
+    neg_pre_grads: list[Tensor],
     *,
     beta1: float,
     beta2: float,
@@ -297,12 +295,12 @@ def _single_tensor_adan(
 
 
 def _multi_tensor_adan(
-    params: List[Tensor],
-    grads: List[Tensor],
-    exp_avgs: List[Tensor],
-    exp_avg_sqs: List[Tensor],
-    exp_avg_diffs: List[Tensor],
-    neg_pre_grads: List[Tensor],
+    params: list[Tensor],
+    grads: list[Tensor],
+    exp_avgs: list[Tensor],
+    exp_avg_sqs: list[Tensor],
+    exp_avg_diffs: list[Tensor],
+    neg_pre_grads: list[Tensor],
     *,
     beta1: float,
     beta2: float,
@@ -356,12 +354,12 @@ def _multi_tensor_adan(
 
 
 def _fused_adan_multi_tensor(
-    params: List[Tensor],
-    grads: List[Tensor],
-    exp_avgs: List[Tensor],
-    exp_avg_sqs: List[Tensor],
-    exp_avg_diffs: List[Tensor],
-    neg_pre_grads: List[Tensor],
+    params: list[Tensor],
+    grads: list[Tensor],
+    exp_avgs: list[Tensor],
+    exp_avg_sqs: list[Tensor],
+    exp_avg_diffs: list[Tensor],
+    neg_pre_grads: list[Tensor],
     *,
     beta1: float,
     beta2: float,
@@ -400,12 +398,12 @@ def _fused_adan_multi_tensor(
 
 
 def _fused_adan_single_tensor(
-    params: List[Tensor],
-    grads: List[Tensor],
-    exp_avgs: List[Tensor],
-    exp_avg_sqs: List[Tensor],
-    exp_avg_diffs: List[Tensor],
-    neg_pre_grads: List[Tensor],
+    params: list[Tensor],
+    grads: list[Tensor],
+    exp_avgs: list[Tensor],
+    exp_avg_sqs: list[Tensor],
+    exp_avg_diffs: list[Tensor],
+    neg_pre_grads: list[Tensor],
     *,
     beta1: float,
     beta2: float,

pyproject.toml

@@ -12,12 +12,12 @@ authors = [
 ]
 dependencies = [
     "torch",
-    "geoopt",
+    "geoopt>=0.5.1",
     "numpy",
     "tqdm",
     "cvxpy",
     "scikit-learn==1.5.1",
-    "scipy<=1.9.3", # Ensure scalar_search_wolfe2 import works correctly
+    "geoopt",
     "datasets"
 ]