Added Enrique's fixes.

src/jace/api.py

@@ -19,11 +19,11 @@
 from jace import stages, translator, util
 
 
-__all__ = ["DEFAUL_BACKEND", "JITOptions", "grad", "jacfwd", "jacrev", "jit"]
+__all__ = ["DEFAULT_BACKEND", "JITOptions", "grad", "jacfwd", "jacrev", "jit"]
 
 _P = ParamSpec("_P")
 
-DEFAUL_BACKEND: Final[str] = "cpu"
+DEFAULT_BACKEND: Final[str] = "cpu"
 
 
 class JITOptions(TypedDict, total=False):
@@ -34,11 +34,11 @@ class JITOptions(TypedDict, total=False):
     `jace.jit`. Furthermore, some additional ones might be supported.
 
     Args:
-    backend: Target platform for which DaCe should generate code. Supported values
-        are `'cpu'` or `'gpu'`.
+        backend: Target platform for which DaCe should generate code. Supported values
+            are `'cpu'` or `'gpu'`.
     """
 
-    backend: str
+    backend: Literal["cpu", "gpu"]
 
 
 @overload
@@ -87,7 +87,7 @@ def jit(
         raise ValueError(
             f"The following arguments to 'jace.jit' are not supported: {', '.join(not_supported_jit_keys)}."
         )
-    if kwargs.get("backend", DEFAUL_BACKEND).lower() not in {"cpu", "gpu"}:
+    if kwargs.get("backend", DEFAULT_BACKEND).lower() not in {"cpu", "gpu"}:
         raise ValueError(f"The backend '{kwargs['backend']}' is not supported.")
 
     def wrapper(f: Callable[_P, Any]) -> stages.JaCeWrapped[_P]:
@@ -99,7 +99,7 @@ def wrapper(f: Callable[_P, Any]) -> stages.JaCeWrapped[_P]:
                 else primitive_translators
             ),
             jit_options=kwargs,
-            device=util.parse_backend_jit_option(kwargs.get("backend", DEFAUL_BACKEND)),
+            device=util.to_device_type(kwargs.get("backend", DEFAULT_BACKEND)),
         )
         functools.update_wrapper(jace_wrapper, f)
         return jace_wrapper

src/jace/tracing.py

@@ -87,10 +87,10 @@ def make_jaxpr(
 
     # NOTE: In the current implementation we are using `jax.make_jaxpr()`. But this
     #   is a different implementation than `jax.jit()` uses. The main difference
-    #   between the two, seems to be the set of arguments that are supported. In JaCe,
+    #   between the two seems to be the set of arguments that are supported. In JaCe,
     #   however, we want to support all arguments that `jace.jit()` does.
     #   For establishing compatibility we have to clear the arguments to make them
-    #   compatible, with what `jax.make_jaxpr()` and `jace.jit()` supports.
+    #   compatible with what `jax.make_jaxpr()` and `jace.jit()` supports.
     trace_options = {}
 
     def tracer_impl(

src/jace/util/__init__.py

@@ -18,8 +18,8 @@
     get_jax_var_shape,
     is_tracing_ongoing,
     move_into_jax_array,
-    parse_backend_jit_option,
     propose_jax_name,
+    to_device_type,
     translate_dtype,
 )
 from .traits import (
@@ -53,7 +53,7 @@
     "is_scalar",
     "is_tracing_ongoing",
     "move_into_jax_array",
-    "parse_backend_jit_option",
     "propose_jax_name",
+    "to_device_type",
     "translate_dtype",
 ]

src/jace/util/jax_helper.py

@@ -30,12 +30,6 @@
     import numpy as np
 
 
-try:
-    import cupy as cp  # type: ignore[import-not-found]
-except ImportError:
-    cp = None
-
-
 @dataclasses.dataclass(repr=True, frozen=True, eq=False)
 class JaCeVar:
     """
@@ -266,7 +260,7 @@ def get_jax_literal_value(lit: jax_core.Atom) -> bool | float | int | np.generic
     raise TypeError(f"Failed to extract value from '{lit}' ('{val}' type: {type(val).__name__}).")
 
 
-def parse_backend_jit_option(
+def to_device_type(
     backend: str | dace.DeviceType,
 ) -> dace.DeviceType:
     """Turn JAX' `backend` option into the proper DaCe device type."""