Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[NFC] Export jacobian_default_mode #1473

Merged
merged 1 commit into from May 8, 2023

Conversation

PhilipVinc
Copy link
Member

@PhilipVinc PhilipVinc commented May 8, 2023

Move and renames the internal chose_qgt from the qgt module to nk.jax.jacobian_default_mode, and properly documents it.

It turns out it's useful to implement NTK and other objects, and GPSKet uses it as well, so let's make it public API together with Jacobian..

This changes nothing, simply moves a function (adds a deprecation warning for the internal function in case anybody uses this function) and adds more extensive docstrings.

@PhilipVinc
Copy link
Member Author

cc @yannra @maxbortone I see you're using it.
I'm moving it to somewhere more relevant, and were it can be kept stable

@codecov
Copy link

codecov bot commented May 8, 2023

Codecov Report

Merging #1473 (b78e9fe) into master (b970084) will decrease coverage by 1.23%.
The diff coverage is 86.27%.

@@            Coverage Diff             @@
##           master    #1473      +/-   ##
==========================================
- Coverage   83.74%   82.52%   -1.23%     
==========================================
  Files         238      239       +1     
  Lines       13667    13456     -211     
  Branches     2075     2075              
==========================================
- Hits        11445    11104     -341     
- Misses       1708     1821     +113     
- Partials      514      531      +17     
Impacted Files Coverage Δ
netket/optimizer/qgt/qgt_jacobian_common.py 94.44% <50.00%> (+1.38%) ⬆️
netket/jax/_jacobian/default_mode.py 88.37% <88.37%> (ø)
netket/jax/__init__.py 100.00% <100.00%> (ø)
netket/jax/_jacobian/__init__.py 100.00% <100.00%> (ø)
netket/optimizer/qgt/qgt_jacobian_dense.py 97.91% <100.00%> (ø)
netket/optimizer/qgt/qgt_jacobian_pytree.py 99.03% <100.00%> (ø)

... and 102 files with indirect coverage changes

@PhilipVinc PhilipVinc merged commit 6fef3b6 into netket:master May 8, 2023
8 of 9 checks passed
@PhilipVinc PhilipVinc deleted the pv/jacobian-mode branch May 8, 2023 14:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

1 participant