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

test: Added an Extensive Set of Tests #21

Open
wants to merge 512 commits into
base: main
Choose a base branch
from
Open
Changes from 1 commit
Commits
Show all changes
512 commits
Select commit Hold shift + click to select a range
345d1bb
Added copy (`copy` and `device_put`) translators, however, there is n…
philip-paul-mueller May 24, 2024
979f4ec
Made some small modifications.
philip-paul-mueller May 27, 2024
2d2c160
Reloacated some code.
philip-paul-mueller May 27, 2024
8adae42
Added a transformation for `slice`.
philip-paul-mueller May 27, 2024
011a2d6
Subclasses of mapped base now have to include the assignment to retur…
philip-paul-mueller May 27, 2024
a3ac868
Implemented a test for teh `select_n` primitive.
philip-paul-mueller May 27, 2024
5f186f6
First, chunk of work.
philip-paul-mueller May 27, 2024
df82038
Small corrections.
philip-paul-mueller May 28, 2024
60b7acb
Size corrections.
philip-paul-mueller May 28, 2024
2190faf
Moved the description of the constructor arguments from `__init__` to…
philip-paul-mueller May 28, 2024
5d6962a
Moved the `jax` package.
philip-paul-mueller May 28, 2024
6c980a9
Fixed the type annotation.
philip-paul-mueller May 28, 2024
6dd65f3
Added a todo.
philip-paul-mueller May 28, 2024
a0996eb
Merge branch 'initial_implementation' into initial_implementation_cont
philip-paul-mueller May 28, 2024
f7fe13d
Found something interesting in DaCe, but it does not fully match the …
philip-paul-mueller May 28, 2024
7b7b7ef
WIP: Dynamic slicing, it seems to work, but it is not nice.
philip-paul-mueller May 28, 2024
28c77ca
Tidied up the dynamic slicing, however, it is not yet fully compatible.
philip-paul-mueller May 29, 2024
054bb22
Dynamic slice now also supports the shifting of windows.
philip-paul-mueller May 29, 2024
c7e8b9f
Some cleaning.
philip-paul-mueller May 29, 2024
a28a9c8
Removed the warning about the useless casting in the tests.
philip-paul-mueller May 29, 2024
5bb4b5e
Started to reorganize the tests.
philip-paul-mueller May 29, 2024
959849a
WIP: Started to better organize the tests.
philip-paul-mueller May 29, 2024
047bb46
Applied a first version to the new code.
philip-paul-mueller May 29, 2024
64135b9
Second round.
philip-paul-mueller May 29, 2024
5756d1e
Third rounds for updated configuration.
philip-paul-mueller May 30, 2024
8b66af8
Fourth batch of adapting new configuration.
philip-paul-mueller May 30, 2024
a5826e9
Fixed a bug in the cache.
philip-paul-mueller May 30, 2024
6f4c45b
Fixed a bug in the cache.
philip-paul-mueller May 30, 2024
4735092
Reorganized some stuff.
philip-paul-mueller May 30, 2024
758adc8
Improved coverage.
philip-paul-mueller May 30, 2024
7fa03c3
chore: add missing dependencies in the mypy hook for pre-commit
egparedes May 30, 2024
ba854e0
remove typing-extensions
egparedes May 30, 2024
6b634bd
Merge remote-tracking branch 'github/main' into initial_implementatio…
philip-paul-mueller May 30, 2024
cd7fc01
Merge remote-tracking branch 'enrique/fix_mypy_deps' into initial_imp…
philip-paul-mueller May 30, 2024
2cdd3fd
First batch of fixes.
philip-paul-mueller May 30, 2024
e4bd289
Merge remote-tracking branch 'github/main' into initial_implementation
philip-paul-mueller May 30, 2024
37c933c
Integrated some of Enrique's changes.
philip-paul-mueller May 30, 2024
699b93e
Rereverted Some of Enrique's suggestions.
philip-paul-mueller May 30, 2024
53aae7b
Made the renaming from driver to builder, that is in my view pointless.
philip-paul-mueller May 30, 2024
eabda1f
Merge branch 'initial_implementation' into initial_implementation_cont
philip-paul-mueller May 30, 2024
f0ea5b4
Made some small changes.
philip-paul-mueller May 30, 2024
fae3ce3
Last screening.
philip-paul-mueller May 31, 2024
71eae32
Merge remote-tracking branch 'github/main' into initial_implementation
philip-paul-mueller May 31, 2024
43ba073
Merge branch 'initial_implementation' into initial_implementation_cont
philip-paul-mueller May 31, 2024
a23f0ab
Started to restructure the tests.
philip-paul-mueller May 31, 2024
555e815
Screened some primitive translators.
philip-paul-mueller May 31, 2024
fc91156
Added more stuff to the ALU translator and discovered something about…
philip-paul-mueller May 31, 2024
1be3237
Because of a bug in the `logical_not` implementation of teh ALU it do…
philip-paul-mueller May 31, 2024
9ab2e9a
Fixed a bug in the `select_n` translator.
Jun 1, 2024
e4a099d
Updated the `mkarray()` function of the tests, it is now able to gene…
Jun 1, 2024
8bdea82
Compacted the tests for select.
Jun 1, 2024
f5eb441
Updated the coverage configuration, however, not yet tested.
Jun 1, 2024
4179d40
Fixed the `TranslatorContext` and `TransltedJaxprSDFG`.
philip-paul-mueller Jun 2, 2024
acfcf39
Now the translation conext and the translated sdfg no longer share an…
philip-paul-mueller Jun 2, 2024
beeb9fd
Merge branch 'initial_implementation' into initial_implementation_cont
philip-paul-mueller Jun 2, 2024
cf99b02
Updated the coverage again.
philip-paul-mueller Jun 2, 2024
7424687
Made a note about what tests are still needed.
philip-paul-mueller Jun 2, 2024
b996041
Fixed the issue with the logical operations.
philip-paul-mueller Jun 2, 2024
bb9d81d
Updated the ALU tests.
philip-paul-mueller Jun 2, 2024
95b34cf
Updated the iota test.
philip-paul-mueller Jun 2, 2024
72c055a
Fixed an issue in the pyproject configuration regarding the `TYPE_CHE…
Jun 2, 2024
66fa038
Fixed a bug in the builder.
philip-paul-mueller Jun 3, 2024
fa599c3
Added some more tests to the empty case.
philip-paul-mueller Jun 3, 2024
e164d75
Added some more tests to the builder.
philip-paul-mueller Jun 3, 2024
7fe6f3b
Fixed a bug in the builder.
philip-paul-mueller Jun 3, 2024
b518ec7
Reenabled the input output functionality in the runner function again.
philip-paul-mueller Jun 3, 2024
94bdf57
Updated some tests, especially the x64 test.
philip-paul-mueller Jun 3, 2024
380dca9
Reenabled the input output functionality in the runner function again.
philip-paul-mueller Jun 3, 2024
35dc1f1
Enrique's comments.
philip-paul-mueller Jun 4, 2024
93e7935
Merge branch 'initial_implementation' into initial_implementation_cont
philip-paul-mueller Jun 4, 2024
ff66995
Modified some tests.
philip-paul-mueller Jun 4, 2024
74435a2
Reorganized some tests.
philip-paul-mueller Jun 4, 2024
acb9501
Updated the calling code.
philip-paul-mueller Jun 4, 2024
5373ca9
It is now also possible to generate random complex.
philip-paul-mueller Jun 4, 2024
e9597f2
Updated the alu tests.
philip-paul-mueller Jun 4, 2024
4499cfb
Enrique's suggestions/comments.
philip-paul-mueller Jun 5, 2024
50b4692
Merge branch 'initial_implementation' into initial_implementation_cont
philip-paul-mueller Jun 5, 2024
5c1e8c6
Updated the tests.
philip-paul-mueller Jun 5, 2024
01cc777
Allied some reformating.
philip-paul-mueller Jun 6, 2024
9ab7360
Merge branch 'initial_implementation' into initial_implementation_cont
philip-paul-mueller Jun 6, 2024
c8b9763
Reapplied some stuff.
philip-paul-mueller Jun 6, 2024
f5b7ccc
Adapted formating in the tests, however, they are so ugly.
philip-paul-mueller Jun 6, 2024
1c3d3a6
Updated the ALT tests.
philip-paul-mueller Jun 6, 2024
7990569
Updated the ALT tests.
philip-paul-mueller Jun 6, 2024
972c9c0
Updated the tests a little bit.
philip-paul-mueller Jun 7, 2024
95f04be
Updated how `is_tracing_ongoing()` works.
philip-paul-mueller Jun 7, 2024
bbc51d0
The translator is now also able to handle cases with no input arguments.
philip-paul-mueller Jun 7, 2024
f1846eb
Updated the tests.
philip-paul-mueller Jun 7, 2024
bf68021
Updated the element type conversion translator.
philip-paul-mueller Jun 7, 2024
a376aad
Updated the slicing primitive and tests.
philip-paul-mueller Jun 7, 2024
8a79f83
It is now also possible that Jax arrays can be used.
philip-paul-mueller Jun 7, 2024
a4697be
Removed some direct import.
philip-paul-mueller Jun 7, 2024
b64b666
Implemented a way to modify the set of compilation/optimization optio…
philip-paul-mueller Jun 7, 2024
ed23d49
Added a fixture that allows to select which options are used.
philip-paul-mueller Jun 7, 2024
3518bdc
Dynamic slicing can not run with disabled optimization.
philip-paul-mueller Jun 7, 2024
7d6cc9c
For the time being disabled the optimized tests in the primitive tran…
philip-paul-mueller Jun 7, 2024
630fcce
First stepe, in fixing the array/scalar stuff.
philip-paul-mueller Jun 7, 2024
e46637d
The Jaxpr is not also included in the translation context, do I need …
philip-paul-mueller Jun 7, 2024
c7f6cc9
Fixed some missed up formating.
philip-paul-mueller Jun 7, 2024
73f0116
I have now disabled the magic comma in ruff.
philip-paul-mueller Jun 7, 2024
ca7f32a
I have now disabled the magic comma in ruff.
philip-paul-mueller Jun 7, 2024
3c6194a
I have now disabled the magic comma in ruff.
philip-paul-mueller Jun 7, 2024
174bea8
Merge branch 'initial_implementation' into initial_implementation_cont
philip-paul-mueller Jun 7, 2024
4f87173
Merge branch 'initial_implementation_cont' into initial_implementatio…
philip-paul-mueller Jun 7, 2024
c5ebe96
Renamed a file.
Jun 9, 2024
45fce56
Fixed an issue in the tests regarding the direct construction of the …
philip-paul-mueller Jun 10, 2024
a2ee92d
Added a layer to handle all outputs and inputs.
philip-paul-mueller Jun 10, 2024
b7fcc80
Updated the tests.
philip-paul-mueller Jun 10, 2024
ceb46be
Merge branch 'initial_implementation_cont_fixing_scalar' into initial…
philip-paul-mueller Jun 10, 2024
b47284d
Some small changes.
philip-paul-mueller Jun 10, 2024
40f8574
Added a configuration file in the `docs/` folder such that the traili…
philip-paul-mueller Jun 10, 2024
50860bb
Merge branch 'initial_implementation' into initial_implementation_cont
philip-paul-mueller Jun 10, 2024
4f52bad
WIP: requirements manager
egparedes May 29, 2024
9af8a4e
Small changes to the documentation of the translators.
philip-paul-mueller Jun 10, 2024
647c5f7
Implemented pytrees.
philip-paul-mueller Jun 11, 2024
c8c612c
Update tooling configs.
egparedes Jun 11, 2024
4788405
Update hooks versions
egparedes Jun 12, 2024
414c55f
Fix code style in markdown file
egparedes Jun 12, 2024
8cec411
Made some cleaning to the pytree implementation.
philip-paul-mueller Jun 12, 2024
9b40b27
Fine tune rule ignores
egparedes Jun 12, 2024
894b10c
More fixes and additions to ruff config
egparedes Jun 12, 2024
1928b95
Add coverage config from Philip
egparedes Jun 12, 2024
9909f6d
More fine tuning of ruff and pre-commit
egparedes Jun 12, 2024
5f24405
Removed a convenience function.
philip-paul-mueller Jun 12, 2024
cc7f649
Fix wrong place to add ruff configuration
egparedes Jun 12, 2024
03d6d08
Updates from new config
egparedes Jun 12, 2024
0da47fa
Fixes
egparedes Jun 12, 2024
b7a166f
Back to defaults
egparedes Jun 12, 2024
367e7b6
Fix format
egparedes Jun 12, 2024
1e825c3
Applied some of enriques new configuration.
philip-paul-mueller Jun 12, 2024
7607fd9
Merge remote-tracking branch 'enrique/more_configs' into initial_impl…
philip-paul-mueller Jun 12, 2024
b759c3b
Merge branch 'main' into req_sync
egparedes Jun 12, 2024
d123166
Final cleanups and fixes
egparedes Jun 12, 2024
39dd108
Clean up sync script
egparedes Jun 12, 2024
889ec88
Merge branch 'initial_implementation' into initial_implementation_cont
philip-paul-mueller Jun 13, 2024
c1074c4
Fixed an issue with the optimization options.
philip-paul-mueller Jun 13, 2024
9ac17bd
Implemented a `CompiledJaxprSDFG` object.
philip-paul-mueller Jun 13, 2024
6a8e2b5
Made some cleaning.
philip-paul-mueller Jun 13, 2024
c05ee3c
Added a test for the pytree support in JaCe.
philip-paul-mueller Jun 13, 2024
ce4c643
Enabled the optimization mode in the tests again.
philip-paul-mueller Jun 13, 2024
1743374
Disabled the elementconversion stuff, since it takes too long.
philip-paul-mueller Jun 13, 2024
efc27e0
Disabled the optimized unit tests, since there is a simplification bu…
philip-paul-mueller Jun 13, 2024
0221d7b
Updated the tests a little bit.
philip-paul-mueller Jun 13, 2024
fd9fb8e
This should solve some cache issue.
philip-paul-mueller Jun 13, 2024
e43dd53
Renamed the `mkarray()` function to `make_array()`.
philip-paul-mueller Jun 13, 2024
744c2b2
Should cover most of the pep whatever naming.
philip-paul-mueller Jun 13, 2024
a11f111
Merge remote-tracking branch 'github/main' into initial_implementation
philip-paul-mueller Jun 14, 2024
9cd0e01
This should fix the problems in the CI.
philip-paul-mueller Jun 14, 2024
3377f41
Now allowed direct return values even in non empty Jaxpr.
philip-paul-mueller Jun 14, 2024
8065ed2
Merge remote-tracking branch 'fork/req_sync_fix' into initial_impleme…
philip-paul-mueller Jun 14, 2024
272dd00
Bump version.
philip-paul-mueller Jun 14, 2024
3e5408b
Bump version.
philip-paul-mueller Jun 14, 2024
3f1d2ad
Updated the ignore file to also include some JaCe stuff.
philip-paul-mueller Jun 14, 2024
e6730d1
Moved tracing to its own top level module.
philip-paul-mueller Jun 14, 2024
82cf898
Updated the ROADMAP.
philip-paul-mueller Jun 14, 2024
da57db9
Improve developer documentation and add devenv session setup to nox
egparedes Jun 14, 2024
6eef078
Add information message
egparedes Jun 14, 2024
1700885
Merge remote-tracking branch 'enrique/req_sync' into initial_implemen…
philip-paul-mueller Jun 14, 2024
90a25a6
Merge remote-tracking branch 'github/main' into initial_implementation
philip-paul-mueller Jun 17, 2024
61638df
Merge branch 'initial_implementation' into initial_implementation_cont
philip-paul-mueller Jun 17, 2024
6369908
pytest now ignores deprication warnings from numpy, this is needed be…
philip-paul-mueller Jun 17, 2024
afe5cdf
Merge branch 'initial_implementation' into initial_implementation_cont
philip-paul-mueller Jun 17, 2024
0aa1726
Changed the names inside the test to make them conforming to an (in t…
philip-paul-mueller Jun 17, 2024
3424a59
Updated the tracing function.
philip-paul-mueller Jun 17, 2024
d8603d5
Now the arguments of the wrapped functions are type chakable.
philip-paul-mueller Jun 18, 2024
54374d3
Added partial support for return value typing in the stages.
philip-paul-mueller Jun 18, 2024
261d902
Imporved typing in the tracing signature.
philip-paul-mueller Jun 18, 2024
027ae35
The options for Jit are now represented as a `TypedDict`.
philip-paul-mueller Jun 18, 2024
7cdb5f5
Integrated part of teh development branch into the new PR branch.
philip-paul-mueller Jun 18, 2024
21e64a8
Merge branch 'extended_infrastructure' into dev_branch
philip-paul-mueller Jun 18, 2024
9fe9e2d
Updated the translators a little bit.
philip-paul-mueller Jun 19, 2024
2ecb7ee
Added teh concatenation translator.
philip-paul-mueller Jun 19, 2024
0b52d4b
First batch of the first review.
philip-paul-mueller Jun 21, 2024
4eb806f
second batch of first review round.
philip-paul-mueller Jun 21, 2024
744738b
DaCe does not consider dashes as valid names.
philip-paul-mueller Jun 21, 2024
71b6422
Changed `finalize_compilation_options()` to `make_final_compilation_o…
philip-paul-mueller Jun 23, 2024
e3a96ad
Merge branch 'extended_infrastructure' into dev_branch
philip-paul-mueller Jun 23, 2024
bcd941f
It seems that `dace.is_array()` does not considers zero dimensional a…
philip-paul-mueller Jun 23, 2024
4be66b1
Fixed a bug in `is_tracing_ongoing()`.
philip-paul-mueller Jun 23, 2024
617c5c1
Updated the `get_jax_literal_value()` function.
philip-paul-mueller Jun 23, 2024
031b33f
Updated the tests after merging the new PR into the development branch.
philip-paul-mueller Jun 23, 2024
c582669
Fixed a bug in `is_tracing_ongoing()`.
philip-paul-mueller Jun 23, 2024
565c69d
Updated the `get_jax_literal_value()` function.
philip-paul-mueller Jun 23, 2024
0dd9404
Imported some files to silence mypy.
philip-paul-mueller Jun 23, 2024
6b69a45
Fixed a naming error.
philip-paul-mueller Jun 24, 2024
f7ee981
Fixed a naming issue.
philip-paul-mueller Jun 24, 2024
8698397
Specified that a transient is exactly written once.
philip-paul-mueller Jun 24, 2024
f3a65d6
Fixed the noxfile, it now uses the correct Sphinx version.
philip-paul-mueller Jun 24, 2024
969f2f8
The README is now included in the documentation.
philip-paul-mueller Jun 24, 2024
29d09e7
It should be `Note:` instead of `Notes:`.
philip-paul-mueller Jun 24, 2024
d029d13
It should be `Note:` insteed of `Notes:`.
philip-paul-mueller Jun 24, 2024
bb938c7
Merge branch 'extended_infrastructure' into dev_branch
philip-paul-mueller Jun 24, 2024
4f343b3
Started with a document that outlines differences between JaCe and JA…
philip-paul-mueller Jun 24, 2024
72d6771
Added the clamp primitive, that we will need later.
philip-paul-mueller Jun 24, 2024
1766bc3
Modified the test function.
philip-paul-mueller Jun 24, 2024
d3dac22
Added a test for proper broadcasting of ternary operators.
philip-paul-mueller Jun 24, 2024
25f52a9
Added a function that is able to incorporate a context into an other …
philip-paul-mueller Jun 24, 2024
ed8f653
Fixed the concatenate translator.
philip-paul-mueller Jun 24, 2024
f4104b8
Added a translator for the `pjit` primitive.
philip-paul-mueller Jun 24, 2024
95853a3
Added a constructor to `JaCeVar` to construct it from an JAX Atom.
philip-paul-mueller Jun 24, 2024
67e546c
Added the functionality to the `pjit` translator to handle literal in…
philip-paul-mueller Jun 24, 2024
0ec665b
The Jaxpr that is translated is now aviable in the lowered object.
philip-paul-mueller Jun 24, 2024
5bc6220
Added tests for the `pjit` primitives.
philip-paul-mueller Jun 24, 2024
5385a5c
Cleared a bit the tests.
philip-paul-mueller Jun 24, 2024
1296b9d
The function to promote literals to constant is now a free function, …
philip-paul-mueller Jun 24, 2024
7615be3
Why was that thing there?
philip-paul-mueller Jun 24, 2024
70710c1
First implementation of the conditional primitive.
philip-paul-mueller Jun 25, 2024
44a8a80
Addressed Enruiques Suggestions.
philip-paul-mueller Jun 25, 2024
c8a6aa8
Squashed commit of the following:
philip-paul-mueller Jun 25, 2024
819eebb
Made a fix to the tests.
philip-paul-mueller Jun 25, 2024
110b6f4
WIP: Started with a primitive.
philip-paul-mueller Jun 25, 2024
380399a
Updated the list of differences.
philip-paul-mueller Jun 25, 2024
320c528
Added tests for the condition.
philip-paul-mueller Jun 25, 2024
e2ebc4e
Implemented propagation of Memlets in the generated graphs.
philip-paul-mueller Jun 26, 2024
071923c
Added a test for scalar branches.
philip-paul-mueller Jun 27, 2024
0869287
Added Enriques suggestions.
philip-paul-mueller Jul 1, 2024
e7fa770
Adapted the tests.
philip-paul-mueller Jul 1, 2024
aa92f01
Applied Enrique's suggestions.
philip-paul-mueller Jul 1, 2024
36e6b6b
Squashed commit of the following:
philip-paul-mueller Jul 1, 2024
a25ecb7
Added Enrique's suggestions.
philip-paul-mueller Jul 2, 2024
3048019
Updated the tests.
philip-paul-mueller Jul 2, 2024
8932815
Merge remote-tracking branch 'github/main' into dev_branch
philip-paul-mueller Jul 31, 2024
f9ee01a
Made some small modifications.
philip-paul-mueller Sep 13, 2024
fa65ee7
Ported the gather translator.
philip-paul-mueller Sep 13, 2024
b9f9427
Made it possible the the `make_array()` function is able scale automa…
philip-paul-mueller Sep 13, 2024
5a3c87f
Added tests for the gather primitive.
philip-paul-mueller Sep 13, 2024
d6265bc
Added all translators from teh development branch.
philip-paul-mueller Sep 13, 2024
25616ae
Modified the condition primitive.
philip-paul-mueller Sep 13, 2024
24d97fb
Nobody needs that thest in this form.
philip-paul-mueller Sep 13, 2024
9f22bd2
Merge branch 'new_translators' into dev_branch
philip-paul-mueller Sep 13, 2024
411bd7b
Since the bug mention in DaCe issue |1595 was resolved we can enable it.
philip-paul-mueller Sep 13, 2024
c8b7d86
First batch of Enrique's suggestions, but still not done.
philip-paul-mueller Sep 24, 2024
cb600d3
Refactored the gather translator.
philip-paul-mueller Sep 24, 2024
c29fc0d
Some more corrections.
philip-paul-mueller Sep 24, 2024
b0ac641
Merge branch 'new_translators' into dev_branch
philip-paul-mueller Sep 24, 2024
3ee8dad
Made it run again.
philip-paul-mueller Sep 24, 2024
e787ee9
Fixed an error in the tests.
philip-paul-mueller Sep 24, 2024
846a345
Fixed some errors.
philip-paul-mueller Sep 24, 2024
a546932
Merge branch 'new_translators' into new_test_suite
philip-paul-mueller Sep 24, 2024
090c3a2
Updated the `make_array()` function.
philip-paul-mueller Sep 25, 2024
e667854
Disabled the simplify pass again in the tests.
philip-paul-mueller Sep 25, 2024
6fe78b0
Fixed an error in the tests.
philip-paul-mueller Sep 25, 2024
d88752a
The test was not really testing anything, okay if it was working.
philip-paul-mueller Sep 25, 2024
2c7b3c8
Applied Enriques primarly fixes.
philip-paul-mueller Sep 26, 2024
5687f40
Merge branch 'new_translators' into new_test_suite
philip-paul-mueller Sep 26, 2024
cc10a48
Splited a test into multiple one, for better pinpointing.
philip-paul-mueller Sep 26, 2024
18bdee9
Apply suggestions from code review
philip-paul-mueller Sep 26, 2024
3588a77
Merge branch 'new_translators' into new_test_suite
philip-paul-mueller Sep 26, 2024
2169470
The import error that panda is not found is now ignored.
philip-paul-mueller Sep 27, 2024
e17b4eb
Merge branch 'main' into new_test_suite
philip-paul-mueller Oct 1, 2024
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Applied Enriques primarly fixes.
philip-paul-mueller committed Sep 26, 2024
commit 2c7b3c8819daaf74bc721c43a2fe08e2ac51b5ec
21 changes: 2 additions & 19 deletions src/jace/translator/jaxpr_translator_builder.py
Original file line number Diff line number Diff line change
@@ -179,24 +179,6 @@ def append_new_state(
self._ctx.terminal_state = new_state
return new_state

def add_orphan_state(
self,
label: str,
) -> dace.SDFGState:
"""
Add a new orphan state to the SDFG.

The state is not connected to any other state, nor it is the new start state.
Except you know what you are doing you should not use this function and
instead use `self.append_new_state()`.

Args:
label: The name of the state.
"""
if not self.is_allocated():
raise RuntimeError("Builder is not allocated.")
return self._ctx.sdfg.add_state(label=label, is_start_block=False)

@property
def arrays(self) -> Mapping[str, dace_data.Data]:
"""
@@ -520,7 +502,8 @@ def _allocate_translation_ctx(
@property
def _ctx(self) -> TranslationContext:
"""Returns the currently active translation context."""
assert len(self._ctx_stack) != 0, "No context is active."
if not self.is_allocated():
raise RuntimeError("The context is not allocated.")
return self._ctx_stack[-1]

def _clear_translation_ctx(self) -> TranslationContext | None:
Original file line number Diff line number Diff line change
@@ -79,19 +79,19 @@ class LogicalOperationTranslator(mapped_base.MappedOperationTranslatorBase):

Args:
prim_name: The name of the primitive that should be handled.
int_tmpl: The template used for the integer case.
bool_tmpl: The template used for the bool case.
bitwise_tmpl: The template used for the bitwise case.
logical_tmpl: The template used for the logical case.

Note:
Since it does not make sense to single out `not` and keep the other
logical operations in `ArithmeticOperationTranslator` all of them are
handled by this class.
"""

def __init__(self, prim_name: str, int_tmpl: str, bool_tmpl: str) -> None:
def __init__(self, prim_name: str, bitwise_tmpl: str, logical_tmpl: str) -> None:
super().__init__(primitive_name=prim_name)
self._int_tmpl = int_tmpl
self._bool_tmpl = bool_tmpl
self._bitwise_tmpl = bitwise_tmpl
self._logical_tmpl = logical_tmpl

@override
def write_tasklet_code(
@@ -101,8 +101,8 @@ def write_tasklet_code(
eqn: jax_core.JaxprEqn,
) -> str:
if all(util.get_jax_var_dtype(invar) is dace.bool_ for invar in eqn.invars):
return self._bool_tmpl
return self._int_tmpl
return self._logical_tmpl
return self._bitwise_tmpl


# Maps the name of an arithmetic JAX primitive to the code template that is used to
@@ -176,17 +176,29 @@ def write_tasklet_code(
# Maps the name of a logical primitive to the two code templates, first the integer
# case and second the boolean case, that are used to create the body of the mapped
# tasklet. They are used to instantiate the `LogicalOperationTranslator` translators.
_LOGICAL_OPERATION_TEMPLATES: Final[dict[str, tuple[str, str]]] = {
"or": ("__out = (__in0) | (__in1)", "__out = (__in0) or (__in1)"),
"not": ("__out = ~(__in0)", "__out = not (__in0)"),
"and": ("__out = (__in0) & (__in1)", "__out = (__in0) and (__in1)"),
"xor": ("__out = (__in0) ^ (__in1)", "__out = (__in0) != (__in1)"),
_LOGICAL_OPERATION_TEMPLATES: Final[dict[str, dict[str, str]]] = {
"or": {
"bitwise_tmpl": "__out = (__in0) | (__in1)",
"logical_tmpl": "__out = (__in0) or (__in1)",
},
"not": {
"bitwise_tmpl": "__out = ~(__in0)",
"logical_tmpl": "__out = not (__in0)",
},
"and": {
"bitwise_tmpl": "__out = (__in0) & (__in1)",
"logical_tmpl": "__out = (__in0) and (__in1)",
},
"xor": {
"bitwise_tmpl": "__out = (__in0) ^ (__in1)",
"logical_tmpl": "__out = (__in0) != (__in1)",
},
}
# fmt: on


# Instantiate the arithmetic and logical translators from the templates.
for pname, ptmpl in _ARITMETIC_OPERATION_TEMPLATES.items():
translator.register_primitive_translator(ArithmeticOperationTranslator(pname, ptmpl))
for pname, (itmpl, btmpl) in _LOGICAL_OPERATION_TEMPLATES.items():
translator.register_primitive_translator(LogicalOperationTranslator(pname, itmpl, btmpl))
for pname, ptmpl in _LOGICAL_OPERATION_TEMPLATES.items(): # type: ignore[assignment] # Type confusion
translator.register_primitive_translator(LogicalOperationTranslator(pname, **ptmpl)) # type: ignore[arg-type] # Type confusion
2 changes: 1 addition & 1 deletion src/jace/translator/primitive_translators/conditions.py
Original file line number Diff line number Diff line change
@@ -116,7 +116,7 @@ def condition_translator(
branch_states.append(branch_state)

# Connect all branch states to the join state
join_state = builder.add_orphan_state(f"{name_pattern}__join_state")
join_state = builder._ctx.sdfg.add_state(label=f"{name_pattern}__join_state")
for branch_state in branch_states:
builder.sdfg.add_edge(
branch_state,
2 changes: 1 addition & 1 deletion src/jace/translator/primitive_translators/slicing.py
Original file line number Diff line number Diff line change
@@ -57,7 +57,7 @@ def make_input_memlets(
eqn: jax_core.JaxprEqn,
) -> dict[str, dace.Memlet]:
strides: Sequence[int] = (
((1,) * len(tskl_ranges)) if eqn.params["strides"] is None else eqn.params["strides"]
eqn.params["strides"] if eqn.params["strides"] else ((1,) * len(tskl_ranges))
)
start_indices: Sequence[int] = eqn.params["start_indices"] # Fist index to slice
return {