Add stricter type checking and expected matrix sizes by corwinjoy · Pull Request #42 · cornellius-gp/linear_operator

corwinjoy · 2022-12-16T01:09:08Z

As per the discussion: cornellius-gp/gpytorch#2180
I have upgraded the library to use Jaxtyping.
This allows a couple things:

The signatures can now include explicit matrix sizes, showing what kinds of results are expected / where broadcasting may be happening.
I have upgraded the operator tests to annotate each tested library with typeguard. So, only during testing, the type hints are tested for correctness to see if the functions actually behave the way the type hints say they will. This helps make the type hints much more accurate so I think its a good step forward.

Issues:

The jaxtyping library requires jax. It should be not to hard to get rid of this dependency (it is on the maintainer's todo list) but that has not been done yet.
It would be nice to be able to add the storage type for the operator (e.g. dense or not) but this again requires upgrading jaxtyping.
I would also like to upgrade the jaxtyping math so I could have expressions like min(M, N) where M,N are array dimensions for some operators.
Because typeguard does not inherit type hint signatures, I've had to copy/paste these from the base _linear_operators class to the derived operators. I think it probably makes sense to create a tool to do this automatically to help with signature maintenance.

…led test cases.

…_operator.py

…rator.py

- Pull requests run a single type checked test (for speed reasons) - Pushes to main and releases run all type checked tests (for completeness)

gpleiss · 2023-05-02T20:35:45Z

Finally got this in! Thanks @corwinjoy !

corwinjoy · 2023-05-02T20:43:32Z

Awesome! Thanks so much for reviewing, polishing and merging this @gpleiss. I know it was a lot to review!

Repeat of cornellius-gp#51, as this snuck in with cornellius-gp#42.

Repeat of #51, as this snuck in with #42.

Balandat · 2023-05-03T15:15:09Z

Hmm @gpleiss I might have missed this, but it looks like jaxtyping is now a required dep even for installation, rather than just testing? In that case we need to add it to install_requires. Will put up a PR for this

See cornellius-gp#42 (comment) ALso removes the `packaging` dep, this is transitive through `pytest`: https://github.com/pytest-dev/pytest/blob/main/setup.cfg#L48

* Make jaxtyping and typeguard required dependencies See #42 (comment) ALso removes the `packaging` dep, this is transitive through `pytest`: https://github.com/pytest-dev/pytest/blob/main/setup.cfg#L48 * Make typeguard an "test" dep only (I don't think it's needed at runtime, is it?) * Remove typeguard dep (transitive via jaxtyping) * Add typeguard restriction to <3.0.0 back in

Balandat · 2023-05-03T16:03:10Z

linear_operator/operators/diag_linear_operator.py


    def _size(self) -> torch.Size:
-        return self._diag.shape + self._diag.shape[-1:]
+        return torch.Size([*self._diag.shape, self._diag.shape[-1]])


This change broke the case where self._diag is a scalar tensor. Will put up a fix

Corwin Joy added 30 commits November 9, 2022 13:31

Copy files from typing to this branch

27fd745

Fix isort config to respect black

c4c5eff

Add jax typing examples

a3575d2

Setup tests to check runtime signatures and begin working through fai…

feb3491

…led test cases.

Continue to fix signatures

d1db0a1

Got runthrough of test cases working after upgrading to latest pytorch

2484f21

Finish signatures for _linear_operator.py

681a8df

Do signatures for added_diag_linear_operator.py

015f4d8

Do signatures for added_diag_linear_operator.py p2

39adfe1

Upgrade signatures and expand scope of runtime tests.

0657a7d

Continue fixing function signature to pass runtime tests.

15b514d

Upgraded signatures and tests for add_diag_linear_operator.py

cc2d1cd

Upgraded signatures and tests for batch_repeat_linear_operator.py

a753ebb

Upgraded signatures and tests for block_diag_linear_operator.py

e4dc417

Upgraded signatures and tests for block_interleaved_linear_operator.py

2f89e5a

Upgraded signatures and tests for block_linear_operator.py

7501082

Upgraded signatures and tests for cat_linear_operator.py.py

2f02a2d

Upgraded signatures and tests for chol_linear_operator.py

10144c4

Upgraded signatures and tests for constant_mul_linear_operator.py

f73a5de

Upgraded signatures and tests for dense_linear_operator.py

9c4399c

Upgraded signatures and tests for identity_linear_operator.py

65c6907

Upgraded signatures and tests for interpolated_linear_operator.py

2d2b8e9

Upgraded signatures and tests for kronecker_product_added_diag_linear…

c900567

…_operator.py

Upgraded signatures and tests for kronecker_product_linear_operator.py

f55a9e9

Upgraded signatures and tests for low_rank_root_added_diag_linear_ope…

92d0181

…rator.py

Upgraded signatures and tests for low_rank_root_linear_operator.py

c339a25

Upgraded signatures and tests for matmul_linear_operator.py

373c421

Upgraded signatures and tests for mul_linear_operator.py

88804b7

Upgraded signatures and tests for psd_sum_linear_operator.py

f33dd58

Upgraded signatures and tests for root_linear_operator.py

7f1a468

gpleiss mentioned this pull request Mar 31, 2023

Format using ufmt, upgrade pre-commit config #56

Merged

gpleiss added 7 commits April 19, 2023 22:32

Merge branch 'main' into jaxtyping

aaf778f

Fix unittests

dac0ce6

Python 3.8 compatability

71346bc

Remove test/typecheck-pytest.ini

80d730a

Fix kronecker product typechecking recursion issues

00afcb6

Fix (unrelated) permutation linear operator failing test

8e19b6c

Fix CatLinearOperator typehint error

8e21e74

gpleiss force-pushed the jaxtyping branch from b95ddee to f012f68 Compare April 24, 2023 16:23

Update CI to run type checking

1411039

- Pull requests run a single type checked test (for speed reasons) - Pushes to main and releases run all type checked tests (for completeness)

gpleiss force-pushed the jaxtyping branch from f012f68 to 1411039 Compare April 24, 2023 16:32

gpleiss added 4 commits April 24, 2023 16:34

Use pytest as a testing dependency

0357ed4

Add pre-commit hook to ensure type hints are propagated

ee64ef9

Fix _diagonal type hint

792ee43

Update contributing guidelines

7491b9b

gpleiss approved these changes May 2, 2023

View reviewed changes

gpleiss added 2 commits May 2, 2023 20:30

Fix some typehints

0d97195

Run a few more tests with type checking

0e1be1a

gpleiss merged commit 32ba847 into cornellius-gp:main May 2, 2023

Balandat added a commit to Balandat/linear_operator that referenced this pull request May 3, 2023

Re-remove @cached decorator for TriangularLinearOperator.to_dense

be544cc

Repeat of cornellius-gp#51, as this snuck in with cornellius-gp#42.

Balandat mentioned this pull request May 3, 2023

Re-remove @cached decorator for TriangularLinearOperator.to_dense #59

Merged

gpleiss pushed a commit that referenced this pull request May 3, 2023

Re-remove @cached decorator for TriangularLinearOperator.to_dense (#59)

c8d4282

Repeat of #51, as this snuck in with #42.

Balandat mentioned this pull request May 3, 2023

Make jaxtyping dep install requirement #60

Merged

Balandat reviewed May 3, 2023

View reviewed changes

Balandat mentioned this pull request Oct 8, 2025

[Proposal] Remove jaxtyping dependency #113

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add stricter type checking and expected matrix sizes#42

Add stricter type checking and expected matrix sizes#42
gpleiss merged 75 commits intocornellius-gp:mainfrom
corwinjoy:jaxtyping

corwinjoy commented Dec 16, 2022

Uh oh!

gpleiss commented May 2, 2023

Uh oh!

corwinjoy commented May 2, 2023

Uh oh!

Balandat commented May 3, 2023

Uh oh!

Balandat May 3, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

corwinjoy commented Dec 16, 2022

Uh oh!

gpleiss commented May 2, 2023

Uh oh!

corwinjoy commented May 2, 2023

Uh oh!

Balandat commented May 3, 2023

Uh oh!

Balandat May 3, 2023

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants