refactor: moved score_matching.py and weights.py to equinox, updated type hints - Githubissues

gchq / coreax

A library for coreset algorithms, written in Jax for fast execution and GPU support.

Apache License 2.0

15 stars 1 forks source link

refactor: moved score_matching.py and weights.py to equinox, updated type hints #657

Closed db091756 closed 2 weeks ago

db091756 commented 2 weeks ago

Closes #131, #656, #417

PR Type

Code style update (formatting, local variables)
Refactoring (no functional changes)

Description

Moves _convert_stein_discrepancy to score_matching.py Refactors score_matching.py and weights.py to use equinox. Fixes various type issues. Test coverage of score_matching.py is 100%

How Has This Been Tested?

Moves _convert_stein_discrepancy unit tests to test_score_matching.py. Adds a small score matching unit test to get to 100% coverage.

Does this PR introduce a breaking change?

(Write your answer here.)

Screenshots

(Write your answer here.)

Checklist before requesting a review

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