Merge pull request #2723 from Balandat/v1152_bump

Remove jaxtyping dep, bump linear_operator req to v0.6.1

Author: Balandat

docs/source/conf.py

@@ -20,8 +20,6 @@
 import sphinx_rtd_theme  # noqa
 import warnings
 
-import jaxtyping
-
 
 def read(*names, **kwargs):
     with io.open(
@@ -262,27 +260,6 @@ def _convert_internal_and_external_class_to_strings(annotation):
     return res
 
 
-# Convert jaxtyping dimensions into strings
-def _dim_to_str(dim):
-    if isinstance(dim, jaxtyping._array_types._NamedVariadicDim):
-        return "..."
-    elif isinstance(dim, jaxtyping._array_types._FixedDim):
-        res = str(dim.size)
-        if dim.broadcastable:
-            res = "#" + res
-        return res
-    elif isinstance(dim, jaxtyping._array_types._SymbolicDim):
-        expr = dim.elem
-        return f"({expr})"
-    elif "jaxtyping" not in str(dim.__class__):  # Probably the case that we have an ellipsis
-        return "..."
-    else:
-        res = str(dim.name)
-        if dim.broadcastable:
-            res = "#" + res
-        return res
-
-
 # Function to format type hints
 def _process(annotation, config):
     """
@@ -295,12 +272,6 @@ def _process(annotation, config):
     if isinstance(annotation, str):
         return annotation
 
-    # Jaxtyping: shaped tensors or linear operator
-    elif hasattr(annotation, "__module__") and "jaxtyping" == annotation.__module__:
-        cls_annotation = _convert_internal_and_external_class_to_strings(annotation.array_type)
-        shape = " x ".join([_dim_to_str(dim) for dim in annotation.dims])
-        return f"{cls_annotation} ({shape})"
-
     # Convert Ellipsis into "..."
     elif annotation == Ellipsis:
         return "..."

setup.py

@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
 
-import io
 import os
 import re
 import sys
@@ -26,7 +25,7 @@
 # Get version
 def find_version(*file_paths):
     try:
-        with io.open(os.path.join(os.path.dirname(__file__), *file_paths), encoding="utf8") as fp:
+        with open(os.path.join(os.path.dirname(__file__), *file_paths), encoding="utf8") as fp:
             version_file = fp.read()
         version_match = re.search(r"^__version__ = version = ['\"]([^'\"]*)['\"]", version_file, re.M)
         return version_match.group(1)
@@ -39,11 +38,10 @@ def find_version(*file_paths):
 
 torch_min = "2.0"
 install_requires = [
-    "jaxtyping",
     "mpmath>=0.19,<=1.3",  # avoid incompatibiltiy with torch+sympy with mpmath 1.4
     "scikit-learn",
     "scipy>=1.6.0",
-    "linear_operator>=0.6",
+    "linear_operator>=0.6.1",
 ]
 # if recent dev version of PyTorch is installed, no need to install stable
 try: