This repository has been archived by the owner on Oct 9, 2024. It is now read-only.
Jax Port of MrVI #1
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
|
I fixed everything but the device at the end in my PRs |
adamgayoso
reviewed
Oct 19, 2022
src/scvi_v2/_model.py
Outdated
Comment on lines
211
to
213
| jit_inference_fn = self.module.get_jit_inference_fn( | ||
| inference_kwargs={"mc_samples": mc_samples, "cf_sample": cf_sample} | ||
| ) |
There was a problem hiding this comment.
I'm pretty sure this will cause the fn to recompile each time, you kind of want a data structure containing each possible inference fn (one per value of sample).
Otherwise you'd want jax to trace over cf_sample, which we can also do
There was a problem hiding this comment.
i think there should be an api here which allows for traced values, (i.e. pass kwargs into the inference_fn, not the get_jit_inference_fn)
adamgayoso
reviewed
Oct 24, 2022
adamgayoso
approved these changes
Oct 24, 2022
Codecov Report
Additional details and impacted files@@ Coverage Diff @@
## main #1 +/- ##
=======================================
Coverage ? 92.89%
=======================================
Files ? 6
Lines ? 394
Branches ? 0
=======================================
Hits ? 366
Misses ? 28
Partials ? 0
Flags with carried forward coverage won't be shown. Click here to find out more. |
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Tests pass using main branch of scvi-tools.
Exact port of torch code with only linear decoder options. Will build off of this.