Feat/jaxtyping #746

db091756 · 2024-08-22T14:27:54Z

PR Type

Refactoring (no functional changes)

Description

I do not touch the typing of the tests in this PR.

Implement Jaxtyping throughout. There should be minimal to no use of ArrayLike or Array without shaping unless the shape is unknown.

There should also be very few Pyright complaints now, except those I couldn't fix with the code as is.

How Has This Been Tested?

Does this PR introduce a breaking change?

(Write your answer here.)

Screenshots

(Write your answer here.)

Checklist before requesting a review

I have made sure that my PR is not a duplicate.
My code follows the style guidelines of this project.
I have commented my code, particularly in hard-to-understand areas.
I have performed a self-review of my code.
I have made corresponding changes to the documentation.
My changes generate no new warnings.
I have added tests that prove my fix is effective or that my feature works.
New and existing unit tests pass locally with my changes.
Any dependent changes have been merged and published in downstream modules.

db091756 · 2024-08-23T14:21:35Z

Performance review

Statistically significant changes

basic_stein:

OLD: compilation 3.193 units ± 0.179 units; execution 2.867 units ± 0.04594 units

NEW: compilation 3.289 units ± 0.2395 units; execution 6.56 units ± 0.03179 units

Significant increase in execution time (128.85%, p=6.059e-29)

Normalisation values for new data: Compilation: 1 unit = 530.66 ms Execution: 1 unit = 681.94 ms

Significant slowdown in the Stein Thinning method, unsure why as I don't believe I have changed anything functionally relevant here

github-actions · 2024-08-23T14:30:41Z

Performance review

Statistically significant changes

basic_stein:
- OLD: compilation 3.193 units ± 0.179 units; execution 2.867 units ± 0.04594 units
- NEW: compilation 3.335 units ± 0.1961 units; execution 6.451 units ± 0.04194 units
- Significant increase in execution time (125.04%, p=1.213e-30)

Normalisation values for new data:
Compilation: 1 unit = 514.29 ms
Execution: 1 unit = 678.56 ms

rg936672

This is an impressive effort! There are a couple of recurring issues:

Nested unions (i.e. Union[A, B, Union[C, D]]) are redundant - they should be flattened out to e.g. Union[A, B, C, D]
Any plan Array type hint that cannot have its shape annotated should have a comment noting why the shape is not possible to statically determine

.cspell/custom_misc.txt

CHANGELOG.md

coreax/approximation.py

coreax/coreset.py

coreax/score_matching.py

coreax/solvers/coresubset.py

coreax/util.py

CHANGELOG.md

rg936672 · 2024-08-28T16:15:38Z

For the remaining pyright complaints (only those in the coreax module, not tests):

The pyright complaints in coreax.kernel can be fixed by replacing each from ... import XYZKernel with from ... import _XYZKernel; the issue seems to just be that the name is reused
The complaint in Coresubset.__init__ can be fixed as described by inheriting from Coreset[Data] rather than Coreset[_Data]
The complaint in networks.create_train_state can be fixed by replacing random_key with jax.random.key(random_key) (or an equivalently aliased import).

With these, the only remaining pyright complaints will be within the tests!

db091756 · 2024-08-29T09:37:53Z

For the remaining pyright complaints (only those in the coreax module, not tests):

The pyright complaints in coreax.kernel can be fixed by replacing each from ... import XYZKernel with from ... import _XYZKernel; the issue seems to just be that the name is reused

The complaint in Coresubset.__init__ can be fixed as described by inheriting from Coreset[Data] rather than Coreset[_Data]

The complaint in networks.create_train_state can be fixed by replacing random_key with jax.random.key(random_key) (or an equivalently aliased import).

With these, the only remaining pyright complaints will be within the tests!

Fixed the last one. Put some questions about the second one in the original comment. The first one I think I am misunderstanding as it seems like this would create import errors. Because the _XYZKernel doesn't exist? I assume you mean from ... import XYZKernel as _XYZKernel, I have implemented that change!

db091756 · 2024-08-29T09:46:19Z

Hopefully responded to all your comments! I also removed some of the overloads that I think are not necessary

github-actions · 2024-08-29T09:49:03Z

Performance review

Statistically significant changes

basic_stein:
- OLD: compilation 3.285 units ± 0.1923 units; execution 2.931 units ± 0.01816 units
- NEW: compilation 3.299 units ± 0.259 units; execution 6.4 units ± 0.048 units
- Significant increase in execution time (118.36%, p=4.212e-22)

Normalisation values for new data:
Compilation: 1 unit = 522.08 ms
Execution: 1 unit = 684.06 ms

db091756 · 2024-08-29T09:53:51Z

Adding more overloads that I missed

github-actions · 2024-08-29T10:16:57Z

Performance review

Statistically significant changes

basic_stein:
- OLD: compilation 3.285 units ± 0.1923 units; execution 2.931 units ± 0.01816 units
- NEW: compilation 3.377 units ± 0.2713 units; execution 6.349 units ± 0.0442 units
- Significant increase in execution time (116.63%, p=4.441e-23)

Normalisation values for new data:
Compilation: 1 unit = 521.73 ms
Execution: 1 unit = 682.68 ms

github-actions · 2024-08-29T10:24:47Z

Performance review

Statistically significant changes

basic_stein:
- OLD: compilation 3.285 units ± 0.1923 units; execution 2.931 units ± 0.01816 units
- NEW: compilation 3.321 units ± 0.2564 units; execution 6.264 units ± 0.05966 units
- Significant increase in execution time (113.72%, p=1.303e-19)

Normalisation values for new data:
Compilation: 1 unit = 512.48 ms
Execution: 1 unit = 686.93 ms

rg936672

Looking good! A few issues remain:

A couple of unresolved comments from the first review
Docs build is failing with py:class reference target not found: 'n *d' and 'n *p'
Tests are now failing!
There are a couple of new Pyright complaints that have popped up (only 8 errors, though). Happy for these to be deferred to a ticket where we include a formal type checker (Include PyRight in CI #427 or Add a runtime type checker #755) if needed though

coreax/kernel.py

rg936672 · 2024-08-30T12:35:10Z

Since other PRs (for example #758) aren't introducing any performance changes, the degradation in Stein performance is an issue created by this PR. This bears investigating!

rg936672 · 2024-08-30T14:31:38Z

Adding the lines

    ("py:class", "'n d'"),
    ("py:class", "'n p'"),

to the Sphinx nitpick_ignore seems to make the docs build pass.

rg936672

Looking good! I've found the source of the performance issues, and the docs build can be fixed by #746 (comment). There's still a couple of Pyright complaints, but I'm happy for these to be left for #427 if needed.

Once the docs build is passing, and the line causing the performance issues is either reverted or given an appropriate TODO and issue, I'm happy for this to be merged.

coreax/score_matching.py

db091756 · 2024-08-30T15:29:39Z

Adding the lines
    ("py:class", "'n d'"),
    ("py:class", "'n p'"),
to the Sphinx nitpick_ignore seems to make the docs build pass.

Ah I was doing
("py:class", "n d"),
("py:class", "n p"),

thanks!

…s JIT wrapped, could be slowing things down

db091756 · 2024-08-30T15:45:39Z

Looking good! I've found the source of the performance issues, and the docs build can be fixed by #746 (comment). There's still a couple of Pyright complaints, but I'm happy for these to be left for #427 if needed.

Once the docs build is passing, and the line causing the performance issues is either reverted or given an appropriate TODO and issue, I'm happy for this to be merged.

I think we leave the remaining Pyright complaints for Beartype to help us with

github-actions · 2024-08-30T16:20:49Z

Performance review

Statistically significant changes

basic_stein:
- OLD: compilation 3.138 units ± 0.2318 units; execution 2.917 units ± 0.03883 units
- NEW: compilation 3.321 units ± 0.2983 units; execution 6.392 units ± 0.05105 units
- Significant increase in execution time (119.12%, p=1.106e-28)

Normalisation values for new data:
Compilation: 1 unit = 528.56 ms
Execution: 1 unit = 682.89 ms

github-actions · 2024-08-30T16:42:34Z

Performance review

Statistically significant changes

basic_stein:
- OLD: compilation 3.138 units ± 0.2318 units; execution 2.917 units ± 0.03883 units
- NEW: compilation 3.214 units ± 0.2804 units; execution 6.386 units ± 0.05752 units
- Significant increase in execution time (118.94%, p=1.077e-26)

Normalisation values for new data:
Compilation: 1 unit = 514.05 ms
Execution: 1 unit = 680.67 ms

github-actions · 2024-08-30T16:59:37Z

Performance review

Statistically significant changes

basic_stein:
- OLD: compilation 3.138 units ± 0.2318 units; execution 2.917 units ± 0.03883 units
- NEW: compilation 3.237 units ± 0.2076 units; execution 6.3 units ± 0.01836 units
- Significant increase in execution time (115.96%, p=4.99e-25)

Normalisation values for new data:
Compilation: 1 unit = 527.58 ms
Execution: 1 unit = 679.46 ms

…loads for dimension handler

github-actions · 2024-09-03T12:21:19Z