Created experimental folder and added all NTK sketching codes #142
Conversation
|
Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA). For more information, open the CLA check for this pull request. |
Codecov Report
@@ Coverage Diff @@
## main #142 +/- ##
==========================================
+ Coverage 86.08% 87.84% +1.76%
==========================================
Files 19 24 +5
Lines 4851 5868 +1017
==========================================
+ Hits 4176 5155 +979
- Misses 675 713 +38
Continue to review full report at Codecov.
|
There was a problem hiding this comment.
Thank you for implementing these features, so far on small tests, the functionality looks good! Here are some small comments to start off the reviews.
Looking forward to see polysketch working, since so far 'rf' conv features don't approximate exact kernel as closely as desired.
|
Hello Amir! Thanks for committing the polysketch routines! Could you sign the contributor license agreement (CLA) as described below in failed checks? Google's github repo requires all the contributors to sign the CLA. |
| kappa0_mat = kappa0(nngp_feat_2d) | ||
| ntk_feat = cholesky(ntk * kappa0_mat).reshape(input_shape + (-1,)) | ||
|
|
||
| else: |
There was a problem hiding this comment.
suggest having the same clause in init_fn above
|
Also, I suggest regularly syncing the fork (https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork) to make sure the PR stays in sync with latest changes. |
There was a problem hiding this comment.
Left some more minor comments that would help us more easily sync this PR with internal codebase. Where possible, please try to to have https://github.com/google/neural-tangents/actions/workflows/pytype.yml pass, since this helps our sync process as well (but feel free to add pytype: disable=<error-code>, or ignore it altogether if the error messages don't make sense!).
…u can run `build_cleaner third_party/py/neural_tangents/...` in the new copybara workspace to create BUILD targets. There's still a lot of stuff to improve (some changes suggested in #142, and here would be nice to run build_cleaner automatically), but this could be a bit helpful in the meantime. PiperOrigin-RevId: 436878167
|
Regarding tests, it looks like Github allows for 20 concurrent jobs running https://docs.github.com/en/actions/learn-github-actions/usage-limits-billing-and-administration#usage-limits, and we currently have 8 (x32/x64 * 4 python versions) linux jobs + 1 pytype job + 2 sketching jobs (x32/x64), so I believe linux/pytype jobs should not slow down the sketching jobs. AFAIK in the last run the sketching job failed first within 25 minutes https://github.com/google/neural-tangents/runs/5845991851?check_suite_focus=true. I'm not sure of the reason, but perhaps reducing the size as I mention in other comments could help? |
| EXACT = 'EXACT' | ||
|
|
||
|
|
||
| def layer(layer_fn): |
There was a problem hiding this comment.
To have a clear picture of the Features API for users, it would be nice to type-annotate its layers and init_fn/feature_fn, by analogy to
and e.g.
(In your case InternalFeatureLayer = Tuple[FeatureInitFn, FeatureFn], and for the functions you can define the typing protocols like
|
Also, the readthedocs fail https://readthedocs.org/projects/neural-tangents/builds/17123387/ should go away after syncing the repo with the main branch |
Created experimental folder and added all NTK sketching codes.