test: Added an Extensive Set of Tests

ROADMAP.md

@@ -7,7 +7,7 @@ A kind of roadmap that gives a rough idea about how the project will be continue
   - [x] Annotation `@jace.jit`.
   - [x] Composable with Jax, i.e. take the Jax derivative of a JaCe annotated function.
   - [x] Implementing the `stages` model that is supported by Jax.
-  - [ ] Handling Jax arrays as native input (only on single host).
+  - [x] Handling Jax arrays as native input (only on single host).
   - [x] Cache the compilation and lowering results for later reuse.
     In Jax these parts (together with the dispatch) are actually written in C++, thus in the beginning we will use a self made cache.
 - [ ] Implementing some basic `PrimitiveTranslators`, that allows us to run some early tests, such as:
@@ -23,6 +23,7 @@ A kind of roadmap that gives a rough idea about how the project will be continue
   But passing these benchmarks could give us some better hint of how to proceed in this matter.
   - [ ] Passing the [pyhpc-benchmark](https://github.com/dionhaefner/pyhpc-benchmarks)
   - [ ] Passing Felix' fluid project; possibility.
+  - [ ] Flash-Attention, there is a DaCe implementation.
 - [ ] Support of static arguments.
 - [ ] Stop relying on `jax.make_jaxpr()`.
   Look at the `jax._src.pjit.make_jit()` function for how to hijack the staging process.

docs/index.md

@@ -3,11 +3,9 @@
 ```{toctree}
 :maxdepth: 2
 :hidden:
-
 ```
 
 ```{include} ../README.md
-:start-after: <!-- SPHINX-START -->
 ```
 
 ## Indices and tables

docs/main_differences.md

@@ -0,0 +1,21 @@
+# Main Differences Between DaCe and JaCe and JAX and JaCe
+
+Essentially JaCe is a frontend that allows DaCe to process JAX code, thus it has to be compatible with both, at least in some sense.
+We will now list the main differences between them, furthermore, you should also consult the ROADMAP.
+
+### JAX vs. JaCe:
+
+- JaCe always traces with enabled `x64` mode.
+  This is a restriction that might be lifted in the future.
+- JAX returns scalars as zero-dimensional arrays, JaCe returns them as array with shape `(1, )`.
+- In JAX parts of the computation runs on CPU parts on GPU, in JaCe everything runs (currently) either on CPU or GPU.
+- Currently JaCe is only able to run on CPU (will be lifted soon).
+- Currently JaCe is not able to run distributed (will be lifted later).
+- Currently not all primitives are supported.
+- JaCe does not return `jax.Array` instances, but NumPy/CuPy arrays.
+- The execution is not asynchronous.
+
+### DaCe vs. JaCe:
+
+- JaCe accepts complex objects using JAX' pytrees.
+- JaCe will support scalar inputs on GPU.

noxfile.py

@@ -134,7 +134,10 @@ def docs(session: nox.Session) -> None:
 @nox.session(reuse_venv=True)
 def api_docs(session: nox.Session) -> None:
     """Build (regenerate) API docs."""
-    session.install(f"sphinx=={REQUIREMENTS['sphinx']}")
+    sphinx_req = REQUIREMENTS["sphinx"]
+    if sphinx_req.isdigit():
+        sphinx_req = "==" + sphinx_req
+    session.install(f"sphinx{sphinx_req}")
     session.chdir("docs")
     session.run(
         "sphinx-apidoc",

pyproject.toml

@@ -113,6 +113,7 @@ addopts = ["-ra", "--showlocals", "--strict-markers", "--strict-config"]
 filterwarnings = [
   "error",
   "ignore:numpy\\..*:DeprecationWarning",  # DaCe is not NumPy v2.0 ready so ignore the usage of deprecated features.
+  "ignore:pandas not found, skipping conversion test\\.:ImportWarning",  # Pandas is not installed on the CI.
 ]
 log_cli_level = "INFO"
 minversion = "6.0"
@@ -233,7 +234,6 @@ max-complexity = 12
 ]
 "tests/**" = [
   "D",  # pydocstyle
-  "N",  # TODO(egparedes): remove ignore as soon as all tests are properly named
   "PLR2004",  # [magic-value-comparison]
   "T10",  # flake8-debugger
   "T20",  # flake8-print

tests/common_fixture.py

@@ -0,0 +1,8 @@
+# JaCe - JAX Just-In-Time compilation using DaCe (Data Centric Parallel Programming)
+#
+# Copyright (c) 2024, ETH Zurich
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Contains all common fixture we need."""

tests/conftest.py

@@ -0,0 +1,102 @@
+# JaCe - JAX Just-In-Time compilation using DaCe (Data Centric Parallel Programming)
+#
+# Copyright (c) 2024, ETH Zurich
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""General configuration for the tests.
+
+Todo:
+    - Implement some fixture that allows to force validation.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Generator
+
+import jax
+import numpy as np
+import pytest
+
+from jace import optimization, stages
+from jace.util import translation_cache as tcache
+
+
+@pytest.fixture(autouse=True)
+def _enable_x64_mode_in_jax() -> Generator[None, None, None]:
+    """Fixture of enable the `x64` mode in JAX.
+
+    Currently, JaCe requires that `x64` mode is enabled and will do all JAX
+    things with it enabled. However, if we use JAX with the intend to compare
+    it against JaCe we must also enable it for JAX.
+    """
+    with jax.experimental.enable_x64():
+        yield
+
+
+@pytest.fixture(autouse=True)
+def _disable_jit() -> Generator[None, None, None]:
+    """Fixture for disable the dynamic jiting in JAX, used by default.
+
+    Using this fixture has two effects.
+    - JAX will not cache the results, i.e. every call to a jitted function will
+        result in a tracing operation.
+    - JAX will not use implicit jit operations, i.e. nested Jaxpr expressions
+        using `pjit` are avoided.
+
+    This essentially disable the `jax.jit` decorator, however, the `jace.jit`
+    decorator is still working.
+
+    Note:
+        The second point, i.e. preventing JAX from running certain things in `pjit`,
+        is the main reason why this fixture is used by default, without it
+        literal substitution is useless and essentially untestable.
+        In certain situation it can be disabled.
+    """
+    with jax.disable_jit(disable=True):
+        yield
+
+
+@pytest.fixture()
+def _enable_jit() -> Generator[None, None, None]:
+    """Fixture to enable jit compilation.
+
+    Essentially it undoes the effects of the `_disable_jit()` fixture.
+    It is important that this fixture is not automatically activated.
+    """
+    with jax.disable_jit(disable=False):
+        yield
+
+
+@pytest.fixture(autouse=True)
+def _clear_translation_cache() -> Generator[None, None, None]:
+    """Decorator that clears the translation cache.
+
+    Ensures that a function finds an empty cache and clears up afterwards.
+    """
+    tcache.clear_translation_cache()
+    yield
+    tcache.clear_translation_cache()
+
+
+@pytest.fixture(autouse=True)
+def _reset_random_seed() -> None:
+    """Fixture for resetting the random seed.
+
+    This ensures that for every test the random seed of NumPy is reset.
+    This seed is used by the `util.mkarray()` helper.
+    """
+    np.random.seed(42)  # noqa: NPY002 [numpy-legacy-random]
+
+
+@pytest.fixture(autouse=True)
+def _set_compile_options() -> Generator[None, None, None]:
+    """Disable all optimizations of jitted code.
+
+    Without explicitly supplied arguments `JaCeLowered.compile()` will not
+    perform any optimizations.
+    Please not that certain tests might override this fixture.
+    """
+    with stages.set_compiler_options(optimization.NO_OPTIMIZATIONS):
+        yield

tests/integration_tests/__init__.py

@@ -0,0 +1,11 @@
+# JaCe - JAX Just-In-Time compilation using DaCe (Data Centric Parallel Programming)
+#
+# Copyright (c) 2024, ETH Zurich
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""JaCe's integration tests.
+
+Currently they are mostly related to the primitive translators.
+"""

tests/integration_tests/primitive_translators/__init__.py

@@ -0,0 +1,8 @@
+# JaCe - JAX Just-In-Time compilation using DaCe (Data Centric Parallel Programming)
+#
+# Copyright (c) 2024, ETH Zurich
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Tests related to the actual primitive subtranslators."""

tests/integration_tests/primitive_translators/conftest.py

@@ -0,0 +1,40 @@
+# JaCe - JAX Just-In-Time compilation using DaCe (Data Centric Parallel Programming)
+#
+# Copyright (c) 2024, ETH Zurich
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""General configuration for the tests of the primitive translators."""
+
+from __future__ import annotations
+
+from collections.abc import Generator
+
+import pytest
+
+from jace import optimization, stages
+
+
+@pytest.fixture(
+    autouse=True,
+    params=[
+        optimization.NO_OPTIMIZATIONS,
+        pytest.param(
+            optimization.DEFAULT_OPTIMIZATIONS,
+            marks=pytest.mark.skip(
+                "Simplify bug 'https://github.com/spcl/dace/issues/1595'; resolved > 16.1"
+            ),
+        ),
+    ],
+)
+def _set_compile_options(request) -> Generator[None, None, None]:
+    """Set the options used for testing the primitive translators.
+
+    This fixture override the global defined fixture.
+
+    Todo:
+        Implement a system that only runs the optimization case in CI.
+    """
+    with stages.set_compiler_options(request.param):
+        yield