feat: Initial implementation #3
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
…rom extern, without the need of having a Jax variable.
…ables. This is basically a simple substitute class that can be used instead of a full jax variable. It is mostly usefull for creating arrays during testing. However, it should also be used to cretae variables for which we do not have anything. This essentially replaces the flags that allowed to explicitly specify shape and dtype in the `add_array()` function.
… single primitive.
But we should split that thing and make it nicer.
The output memlet for the scalar case was not handled correctly.
It is only for debugging.
…one go. This function is highly internal and should not be used.
The function is now able to handle more recent Jax functions, however, it now produces very starnge names. I think we should have some kind of context, similar to the one used by Jax itself.
…n replacement. I think that we will need it in the future, for example it is, in my view, currently the best place to put `jace.jit`. Another idea that could be worthwhile to consider is to mimick the `jace` package after `jax`.
However, it does not cache the compiled code yet. And it is not tested yet.
There is still the problem of scalars as return values. However, even DaCe turns them into arrays.
…ementation This fully introduces Enrique's configuration.
philip-paul-mueller
force-pushed
the
initial_implementation
branch
from
4e2dc0c to
a11f111
Compare
…cause DaCe uses some of these internals.
Welcome to Codecov 🎉Once you merge this PR into your default branch, you're all set! Codecov will compare coverage reports and display results in all future pull requests. Thanks for integrating Codecov - We've got you covered ☂️ |
egparedes
approved these changes
Jun 17, 2024
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.
Initial JaCe code.
This PR introduces a series of very basic functionalities, that will be extended in subsequent PRs.
Most importantly it introduces the basic infrastructure that allows to translate a Jaxpr object into an SDFG.
Furthermore, it also adds the
jace.jitdecorator, which can be used as a replacement forjax.jit.However, a function that was decorated with
jace.jitremains fully composable with Jax transformation, such asjax.gradorjax.jacfwd.But, the functionality is still very basic and
jace.jit, essentially does not accepts any arguments and only works on CPU.Nevertheless, there is a cache for caching tracing, translation and compilation of wrapped functions.
Although, this PR introduces the components for the translation, the actual primitive translators are not yet implemented (this commit adds the
ALUTranslator, but this one was back ported from the prototype for simple tests).As a last point the tests, located in
testsare only a first version, that were not included in the review of this PR.They will be reviewed at a later stage
For more information see the
ROADMAP.mdfile.