Remove jaxtyping dependency

.conda/meta.yaml

@@ -21,7 +21,6 @@ requirements:
 
   run:
     - python>=3.10
-    - jaxtyping
     - mpmath>=0.19,<=1.3
     - pytorch>=2.0
     - scipy

.github/workflows/deploy.yml

@@ -12,7 +12,7 @@ jobs:
     uses: ./.github/workflows/run_linter.yml
 
   run_test_suite:
-    uses: ./.github/workflows/run_type_checked_test_suite.yml
+    uses: ./.github/workflows/run_test_suite.yml
 
   deploy_pypi:
     runs-on: ubuntu-latest

.github/workflows/pull_request.yml

@@ -13,8 +13,3 @@ jobs:
 
   run_test_suite:
     uses: ./.github/workflows/run_test_suite.yml
-
-  run_small_type_checked_test_suite:
-    uses: ./.github/workflows/run_type_checked_test_suite.yml
-    with:
-      files_to_test: "test/operators/test_dense_linear_operator.py test/operators/test_diag_linear_operator.py test/operators/test_kronecker_product_linear_operator.py"

.github/workflows/push_to_main.yml

@@ -12,4 +12,4 @@ jobs:
     uses: ./.github/workflows/run_linter.yml
 
   run_test_suite:
-    uses: ./.github/workflows/run_type_checked_test_suite.yml
+    uses: ./.github/workflows/run_test_suite.yml

.gitignore

@@ -1,3 +1,6 @@
+# Claude
+CLAUDE.md
+
 # Atom plugin files and ctags
 .ftpconfig
 .ftpconfig.cson

.pre-commit-config.yaml

@@ -38,11 +38,3 @@ repos:
     hooks:
     -   id: forbid-crlf
     -   id: forbid-tabs
--   repo: local
-    hooks:
-    -   id: propagate-type-hints
-        name: Propagate Type Hints
-        entry: ./.hooks/propagate_type_hints.sh
-        language: script
-        pass_filenames: true
-        require_serial: true

CONTRIBUTING.md

@@ -39,38 +39,20 @@ We use [standard sphinx docstrings](https://sphinx-rtd-tutorial.readthedocs.io/e
 LinearOperator aims to be fully typed using Python 3.10+
 [type hints](https://www.python.org/dev/peps/pep-0484/).
 We expect any contributions to also use proper type annotations.
-We are using [jaxtyping](https://github.com/google/jaxtyping) to help us be declarative about the dimension sizes used
-in the LinearOperator methods.
-The use of [jaxtyping](https://github.com/google/jaxtyping)  makes it clearer what the functions are doing algebraically
-and where broadcasting is happening.
-
-These type hints are checked in the unit tests by using
-[typeguard](https://github.com/agronholm/typeguard) to perform run-time
-checking of the signatures to make sure they are accurate.
-The signatures are written into the base linear operator class in `_linear_oparator.py`.
-These signatures are then copied to the derived classes by running the script
-`python ./.hooks/propagate_type_hints.py`.
-This is done for:
-1. Consistency. Make sure the derived implementations are following the promised interface.
-2. Visibility. Make it easy to see what the expected signature is, along with dimensions. Repeating the signature in the derived classes enhances readability.
-3. Necessity. The way that jaxtyping and typeguard are written, they won't run type checks unless type annotations are present in the derived method signature.
-
-In short, if you want to update the type hints, update the code in the LinearOperator class in
-`_linear_oparator.py` then run `python ./.hooks/propagate_type_hints.py` to copy the new signature to the derived
-classes.
-
-### Unit Tests
 
-#### With type checking (slower, but more thorough)
-To run the unittests with type checking, run
-```bash
-pytest --jaxtyping-packages=linear_operator,typeguard.typechecked
+Dimension information is documented using inline comments with the format:
+```python
+def _matmul(
+    self: LinearOperator,  # shape: (*batch, M, N)
+    rhs: Tensor,  # shape: (*batch2, N, C) or (*batch2, N)
+) -> Tensor:  # shape: (..., M, C) or (..., M)
 ```
 
-- To run tests within a specific directory, run (e.g.) `pytest test/operators --jaxtyping-packages=linear_operator,typeguard.typechecked`.
-- To run a specific file, run (e.g.) `pytest test/operators/test_matmul_linear_operator.py --jaxtyping-packages=linear_operator,typeguard.typechecked`.
+This convention makes it clear what the functions are doing algebraically
+and where broadcasting is happening.
+
+### Unit Tests
 
-#### Without type checking (faster, but less thorough)
 We use python's `unittest` to run unit tests:
 ```bash
 python -m unittest

docs/source/conf.py

@@ -19,7 +19,6 @@
 import warnings
 from typing import ForwardRef
 
-import jaxtyping
 import sphinx_rtd_theme  # noqa
 
 
@@ -123,27 +122,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):
     """
@@ -156,12 +134,6 @@ def _process(annotation, config):
     if type(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 "..."