Add jax implementation of DII #146
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Proposed changes
Comment 1 - JAX version
The code works with jax=0.4.30, which unfortunately is not available for Python<3.9. Since Python 3.8 is still widely employed I found a workaround by installing this JAX version only for Python>=3.9. For Python<3.9 then the DiffImbalance class is not available (this is specified in the README and in the documentation), and the corresponding tests are not run during the installation.
Comment 2 - JAX on GPU
To avoid installation problems, the JAX library is installed by default on CPU (as from PR #140). The GPU installation depends on several factors, so I would leave it completely to the user (I added some tips in the README and a reference to the installation tutorial by Google).
TODO
The tutorial can be expanded.
The code is missing of a method for the forward / backward greedy selection that is present in the Cython implementation