diff --git a/src/jace/translator/__init__.py b/src/jace/translator/__init__.py index 8ca5476..d6bb0c7 100644 --- a/src/jace/translator/__init__.py +++ b/src/jace/translator/__init__.py @@ -10,12 +10,10 @@ from __future__ import annotations from .jaxpr_translator_driver import JaxprTranslationDriver -from .primitive_translator import PrimitiveTranslator from .translated_jaxpr_sdfg import TranslatedJaxprSDFG __all__ = [ - "PrimitiveTranslator", "JaxprTranslationDriver", "TranslatedJaxprSDFG", ] diff --git a/src/jace/translator/_translation_context.py b/src/jace/translator/_translation_context.py index a0c1854..6d9c43f 100644 --- a/src/jace/translator/_translation_context.py +++ b/src/jace/translator/_translation_context.py @@ -12,9 +12,9 @@ from collections.abc import MutableMapping, Sequence import dace -from jax import core as jcore +from jax import core as jax_core -from jace import translator as jtrans, util as jutil +from jace import translator, util class _TranslationContext: @@ -70,7 +70,7 @@ def __init__( rev_idx: The revision index of the context. name: Name of the SDFG object. """ - if isinstance(name, str) and not jutil._VALID_SDFG_OBJ_NAME.fullmatch(name): + if isinstance(name, str) and not util._VALID_SDFG_OBJ_NAME.fullmatch(name): raise ValueError(f"'{name}' is not a valid SDFG name.") self._sdfg: dace.SDFG = dace.SDFG(name=(name or f"unnamed_SDFG_{id(self)}")) @@ -78,14 +78,14 @@ def __init__( label="initial_state", is_start_block=True ) self._terminal_state: dace.SDFGState = self._start_state - self._jax_name_map: MutableMapping[jcore.Var | jutil.JaCeVar, str] = {} + self._jax_name_map: MutableMapping[jax_core.Var | util.JaCeVar, str] = {} self._inp_names: tuple[str, ...] = () self._out_names: tuple[str, ...] = () self._rev_idx: int = rev_idx - def to_translated_jaxpr_sdfg(self) -> jtrans.TranslatedJaxprSDFG: + def to_translated_jaxpr_sdfg(self) -> translator.TranslatedJaxprSDFG: """Transforms `self` into a `TranslatedJaxprSDFG`.""" - return jtrans.TranslatedJaxprSDFG( + return translator.TranslatedJaxprSDFG( sdfg=self._sdfg, start_state=self._start_state, terminal_state=self._terminal_state, @@ -114,7 +114,7 @@ def terminal_state( self._terminal_state = new_term_state @property - def jax_name_map(self) -> MutableMapping[jcore.Var | jutil.JaCeVar, str]: + def jax_name_map(self) -> MutableMapping[jax_core.Var | util.JaCeVar, str]: return self._jax_name_map @property diff --git a/src/jace/translator/jaxpr_translator_driver.py b/src/jace/translator/jaxpr_translator_driver.py index c6fdb46..5b01fd2 100644 --- a/src/jace/translator/jaxpr_translator_driver.py +++ b/src/jace/translator/jaxpr_translator_driver.py @@ -14,9 +14,10 @@ import dace import jax from dace import data as ddata, properties as dprop -from jax import core as jcore +from jax import core as jax_core -from jace import translator as jtrans, util as jutil +from jace import translator, util +from jace.translator import sub_translators class JaxprTranslationDriver: @@ -81,7 +82,7 @@ def __init__( # They are partitioned by the names of the primitive they have registered for. # This member is allocated by '_init_sub_translators()' and remains allocated # during the lifetime of the object. - self._sub_translators: dict[str, jtrans.PrimitiveTranslator] = None # type: ignore[assignment] + self._sub_translators: dict[str, translator.PrimitiveTranslator] = None # type: ignore[assignment] # These names can not be used for the automatic naming of Jax variables. # They differ from the forbidden names, that they denote valid SDFG names. @@ -104,14 +105,14 @@ def __init__( def translate_jaxpr( self, - jaxpr: jcore.ClosedJaxpr, + jaxpr: jax_core.ClosedJaxpr, *, inp_scalar_as_array: bool = False, name: str | None = None, reserved_names: str | Collection[str] | None = None, allow_empty_jaxpr: bool = False, **kwargs: Any, - ) -> jtrans.TranslatedJaxprSDFG: + ) -> translator.TranslatedJaxprSDFG: """Perform the translation of a Jaxpr into a SDFG. In case this function is called and `self` has an ongoing translation process, a new translation context will be created. @@ -134,7 +135,7 @@ def translate_jaxpr( """ if (len(jaxpr.eqns) == 0) and (not allow_empty_jaxpr): raise ValueError("Passed an empty Jaxpr, but did not allow for empty Jaxpr.") - if not isinstance(jaxpr, jcore.ClosedJaxpr): + if not isinstance(jaxpr, jax_core.ClosedJaxpr): raise TypeError(f"Expected a 'jax.core.ClosedJaxp' instance but got '{type(jaxpr)}'") if len(jaxpr.effects) != 0: raise NotImplementedError("'Jaxpr' with side effects are not supported.") @@ -162,7 +163,7 @@ def translate_jaxpr( jaxpr=jaxpr, inp_scalar_as_array=inp_scalar_as_array, ) - jsdfg: jtrans.TranslatedJaxprSDFG = self._translate_jaxpr_internal(jaxpr) + jsdfg: translator.TranslatedJaxprSDFG = self._translate_jaxpr_internal(jaxpr) # If the translation context is not cleared `self` and `jsdfg` will share the same data. # There is some legitimate use for that. @@ -195,7 +196,7 @@ def append_new_state( prev_state: Alternative `SDFGState` at which we should append the new state. """ - if isinstance(label, str) and (not jutil._VALID_SDFG_OBJ_NAME.fullmatch(label)): + if isinstance(label, str) and (not util._VALID_SDFG_OBJ_NAME.fullmatch(label)): raise ValueError(f"Can not create state with label '{label}' since it is invalid.") # Decide if appending to that state will modify the terminal state. @@ -228,7 +229,7 @@ def get_arrays(self) -> Mapping[str, ddata.Data]: def get_array( self, - name: str | jcore.Atom | jutil.JaCeVar, + name: str | jax_core.Atom | util.JaCeVar, ) -> ddata.Data: """Returns the SDFG `Data` object `name` referees to. @@ -237,7 +238,7 @@ def get_array( """ if isinstance(name, str): sdfg_name: str = name - elif isinstance(name, (jcore.Var, jutil.JaCeVar)): + elif isinstance(name, (jax_core.Var, util.JaCeVar)): sdfg_name = self.map_jax_var_to_sdfg(name) else: raise TypeError(f"Does not know how to handle '{type(name).__name__}'.") @@ -248,19 +249,19 @@ def get_array( @overload def map_jax_var_to_sdfg( self, - jax_var: str | jcore.Atom | jutil.JaCeVar, + jax_var: str | jax_core.Atom | util.JaCeVar, ) -> str: ... @overload def map_jax_var_to_sdfg( self, - jax_var: str | jcore.Atom | jutil.JaCeVar, + jax_var: str | jax_core.Atom | util.JaCeVar, allow_fail: bool, ) -> str | None: ... def map_jax_var_to_sdfg( self, - jax_var: str | jcore.Atom | jutil.JaCeVar, + jax_var: str | jax_core.Atom | util.JaCeVar, allow_fail: bool = False, ) -> str | None: """Get the _name_ of the SDFG variable to which `jax_var` is referring to. @@ -273,7 +274,7 @@ def map_jax_var_to_sdfg( """ if isinstance(jax_var, str): sdfg_name: str = jax_var - elif isinstance(jax_var, jcore.Literal): + elif isinstance(jax_var, jax_core.Literal): raise RuntimeError("There is no SDFG variable for literal '{jax_var}'.") elif jax_var in self._ctx.jax_name_map: sdfg_name = self._ctx.jax_name_map[jax_var] @@ -336,7 +337,7 @@ def get_rev_idx(self) -> int: def add_jax_name_mapping( self, - jax_var: jcore.Var | jutil.JaCeVar, + jax_var: jax_core.Var | util.JaCeVar, sdfg_name: str, ) -> JaxprTranslationDriver: """Creates a mapping between `jax_var` to `sdfg_name`. @@ -384,7 +385,7 @@ def add_reserved_names( raise TypeError(f"Does not know how to handle the type '{type(reserved_names)}'.") for rev_name in reserved_names: assert isinstance(rev_name, str) - if not jutil._VALID_SDFG_VAR_NAME.fullmatch(rev_name): + if not util._VALID_SDFG_VAR_NAME.fullmatch(rev_name): raise ValueError( f"Can not use '{rev_name}' as reserved name as it is not a valid SDFG name." ) @@ -393,7 +394,7 @@ def add_reserved_names( def add_array( self, - arg: jcore.Atom | jutil.JaCeVar, + arg: jax_core.Atom | util.JaCeVar, *, as_transient: bool = True, alt_name: str | None = None, @@ -471,8 +472,8 @@ def add_array( """ assert self.is_allocated() - shape: Sequence[int] = jutil.get_jax_var_shape(arg) - dtype = jutil.get_jax_var_dtype(arg) + shape: Sequence[int] = util.get_jax_var_shape(arg) + dtype = util.get_jax_var_dtype(arg) offset = None # i.e. no offset storage: dace.StorageType = dace.StorageType.Default # Set at later stages (optimization) is_scalar: bool = shape == () @@ -493,7 +494,7 @@ def add_array( raise ValueError( f"Specified 'force_jax_name', but passed '{name_prefix}' as 'name_prefix'." ) - alt_name = jutil._propose_jax_name(arg, self._ctx.jax_name_map) + alt_name = util._propose_jax_name(arg, self._ctx.jax_name_map) if alt_name is not None: assert isinstance( alt_name, str @@ -503,7 +504,7 @@ def add_array( raise ValueError("Passed an empty 'alt_name'.") if alt_name in self._forbidden_names: raise ValueError("'alt_name' is a forbidden name.") - if not jutil._VALID_SDFG_VAR_NAME.fullmatch(alt_name): + if not util._VALID_SDFG_VAR_NAME.fullmatch(alt_name): raise ValueError(f"The passed name 'alt_name' '{alt_name}' is invalid.") if name_prefix is not None: raise ValueError( @@ -536,8 +537,8 @@ def add_array( # Depending on the situation, we will further manipulate it. if alt_name is not None: prop_name = alt_name # Just for completion: will be ignored later - elif isinstance(arg, (jcore.Var, jutil.JaCeVar)): - prop_name = jutil._propose_jax_name(arg, self._ctx.jax_name_map) + elif isinstance(arg, (jax_core.Var, util.JaCeVar)): + prop_name = util._propose_jax_name(arg, self._ctx.jax_name_map) if prop_name.startswith("__"): raise ValueError( f"You tried to create the variable '{prop_name}' which" @@ -545,7 +546,7 @@ def add_array( ) if name_prefix is not None: prop_name = name_prefix + prop_name - elif isinstance(arg, jcore.Literal): # type: ignore[unreachable] + elif isinstance(arg, jax_core.Literal): # type: ignore[unreachable] if not allow_literals: raise NotImplementedError("Jax Literals are not supported.") if alt_name is None: @@ -590,7 +591,7 @@ def add_array( raise ValueError(f"Can't create variable '{arg_name}', name is forbidden.") if arg_name in self._ctx.sdfg.arrays: raise ValueError(f"Can't create variable '{arg_name}', variable is already created.") - if not jutil._VALID_SDFG_VAR_NAME.fullmatch(arg_name): + if not util._VALID_SDFG_VAR_NAME.fullmatch(arg_name): raise ValueError(f"The requested variable name '{arg_name}' is invalid.") # Promotion of scalar to array. @@ -642,7 +643,7 @@ def add_array( def create_jax_var_list( self, - jax_var_list: Sequence[jcore.Atom | jutil.JaCeVar], + jax_var_list: Sequence[jax_core.Atom | util.JaCeVar], prevent_creation: bool = False, only_creation: bool = False, handle_literals: bool = False, @@ -677,11 +678,11 @@ def create_jax_var_list( ret_list: list[None | str] = [] for jax_var in jax_var_list: - if isinstance(jax_var, jcore.Literal): + if isinstance(jax_var, jax_core.Literal): if not handle_literals: raise ValueError("Encountered a literal but `handle_literals` was `False`.") sdfg_name = None - elif isinstance(jax_var, (jcore.Var, jutil.JaCeVar)): + elif isinstance(jax_var, (jax_core.Var, util.JaCeVar)): mapped_sdfg_name: str | None = self.map_jax_var_to_sdfg(jax_var, allow_fail=True) if (mapped_sdfg_name is None) and prevent_creation: raise ValueError(f"'prevent_creation' given but have to create '{jax_var}'.") @@ -702,7 +703,7 @@ def create_jax_var_list( def _create_initial_input( self, - jaxpr: jcore.ClosedJaxpr, + jaxpr: jax_core.ClosedJaxpr, inp_scalar_as_array: bool, ) -> Sequence[str]: """This function will create the internal input variables that are used for the SDFG. @@ -744,7 +745,7 @@ def _create_initial_input( def _create_constants( self, - jaxpr: jcore.ClosedJaxpr, + jaxpr: jax_core.ClosedJaxpr, ) -> Sequence[str]: """Creates all constants requested by the `jaxpr`. @@ -825,21 +826,21 @@ def _init_sub_translators( The function forwards `kwargs` to the constructor of the subtranslators. However, it will remove all arguments starting with an underscore. """ - from jace.translator.sub_translators import _get_subtranslators_cls # Avoid import cycle - assert self._sub_translators is None subtrans_args = {k: v for k, v in subtrans_args.items() if not k.startswith("_")} # type: ignore[unreachable] - sub_translators: dict[str, jtrans.PrimitiveTranslator] = {} - for sub_translator_cls in _get_subtranslators_cls(): - sub_translator: jtrans.PrimitiveTranslator = sub_translator_cls.CREATE(**subtrans_args) - handled_primitives: Iterable[str] = jutil.as_sequence(sub_translator.primitive) + prim_translators: dict[str, translator.PrimitiveTranslator] = {} + for prim_translator_cls in sub_translators._get_subtranslators_cls(): + prim_translator: translator.PrimitiveTranslator = prim_translator_cls.CREATE( + **subtrans_args + ) + handled_primitives: Iterable[str] = util.as_sequence(prim_translator.primitive) for handled_primitive in handled_primitives: - if handled_primitive in sub_translators: - raise RuntimeError(f"Multiple sub_translators for '{handled_primitive}' found.") - sub_translators[handled_primitive] = sub_translator - self._sub_translators = sub_translators + if handled_primitive in prim_translators: + raise RuntimeError(f"Multiple sub translators for '{handled_primitive}' found.") + prim_translators[handled_primitive] = prim_translator + self._sub_translators = prim_translators return self @@ -872,8 +873,8 @@ def _clear_translation_ctx(self) -> JaxprTranslationDriver: def _find_sub_translator_for( self, - eqn: jcore.JaxprEqn, - ) -> jtrans.PrimitiveTranslator: + eqn: jax_core.JaxprEqn, + ) -> translator.PrimitiveTranslator: """Returns the appropriate subtranslator for equation `eqn`.""" assert self._sub_translators is not None @@ -885,8 +886,8 @@ def _find_sub_translator_for( def _translate_single_eqn( self, - jaxpr: jcore.ClosedJaxpr, - eqn: jcore.JaxprEqn, + jaxpr: jax_core.ClosedJaxpr, + eqn: jax_core.JaxprEqn, ) -> tuple[Sequence[str | None], Sequence[str]]: """Translate `eqn` into its SDFG equivalent. @@ -904,8 +905,8 @@ def _translate_single_eqn( While `jaxpr` must be a `ClosedJaxpr`, `eqn` must come from the unclosed instance. The function will perform some consistency checking after the subtranslator was called. """ - assert isinstance(eqn, jcore.JaxprEqn) - assert isinstance(jaxpr, jcore.ClosedJaxpr) + assert isinstance(eqn, jax_core.JaxprEqn) + assert isinstance(jaxpr, jax_core.ClosedJaxpr) if len(eqn.effects) != 0: raise NotImplementedError(f"Equation '{eqn}' has side effects.") @@ -925,7 +926,7 @@ def _translate_single_eqn( ) # Find the subtranslator - subtranslator: jtrans.PrimitiveTranslator = self._find_sub_translator_for(eqn) + subtranslator: translator.PrimitiveTranslator = self._find_sub_translator_for(eqn) # Create the state into which the equation should be translated last_term_state: dace.SDFGState = self.get_terminal_sdfg_state() # noqa: F841 # Will be used later @@ -963,7 +964,7 @@ def _translate_single_eqn( ) for expectedSDFGName, jax_var in zip(out_var_names, eqn.outvars, strict=True): mapped_sdfg_name = self.map_jax_var_to_sdfg(jax_var) - jax_name = jutil.get_jax_var_name(jax_var) + jax_name = util.get_jax_var_name(jax_var) if mapped_sdfg_name != expectedSDFGName: raise ValueError( f"Mapping inconsistency detected, expected that Jax variable" @@ -980,13 +981,13 @@ def _translate_single_eqn( pass elif isinstance(sdfg_var, dace.data.View): raise TypeError( - f"For Jax variable '{jutil.get_jax_var_name(jax_var)}' (SDFG: '{outVarName}')," + f"For Jax variable '{util.get_jax_var_name(jax_var)}' (SDFG: '{outVarName}')," f" which is an output, you used a View, which is not possible." " It must either be an array or a scalar." ) else: raise NotImplementedError( - f"Output variable '{jutil.get_jax_var_name(jax_var)}' (SDFG: '{outVarName}')" + f"Output variable '{util.get_jax_var_name(jax_var)}' (SDFG: '{outVarName}')" f" is of type '{type(sdfg_var).__name__}' which I does not know how to handle." ) @@ -997,8 +998,8 @@ def _translate_single_eqn( def _translate_jaxpr_internal( self, - jaxpr: jcore.ClosedJaxpr, - ) -> jtrans.TranslatedJaxprSDFG: + jaxpr: jax_core.ClosedJaxpr, + ) -> translator.TranslatedJaxprSDFG: """Performs the actual translation of the Jaxpr into an SDFG. The function assumes that the context is allocated as well as initial variables. @@ -1014,7 +1015,7 @@ def _translate_jaxpr_internal( this is used by Jax to indicate that they are never read. Such variables are included by some transformations such as `grad()`. """ - assert isinstance(jaxpr, jcore.ClosedJaxpr) + assert isinstance(jaxpr, jax_core.ClosedJaxpr) assert self.is_allocated() nb_translated_eqn: int = 0 @@ -1023,9 +1024,9 @@ def _translate_jaxpr_internal( assert len(eqn.effects) == 0 if len(eqn.outvars) == 0: # Do we need this special case. continue # Looks more like internal Jax error. - if any(jutil.is_drop_var(outVar) for outVar in eqn.outvars): + if any(util.is_drop_var(outVar) for outVar in eqn.outvars): assert (len(eqn.outvars) == 1) or all( - jutil.is_drop_var(outVar) for outVar in eqn.outvars + util.is_drop_var(outVar) for outVar in eqn.outvars ) continue _, out_var_names = self._translate_single_eqn(jaxpr=jaxpr, eqn=eqn) @@ -1038,7 +1039,7 @@ def _translate_jaxpr_internal( return self._export_context() - def _export_context(self) -> jtrans.TranslatedJaxprSDFG: + def _export_context(self) -> translator.TranslatedJaxprSDFG: """Encapsulate the translation context of `self` into a `TranslatedJaxprSDFG` object.. This function will not deallocate the internal context of `self`. @@ -1049,7 +1050,7 @@ def _export_context(self) -> jtrans.TranslatedJaxprSDFG: assert all((isinstance(x, str) and (len(x) > 0)) for x in self._ctx.inp_names) assert all((isinstance(x, str) and (len(x) > 0)) for x in self._ctx.out_names) - return jtrans.TranslatedJaxprSDFG( + return translator.TranslatedJaxprSDFG( sdfg=self._ctx.sdfg, start_state=self._ctx.start_state, terminal_state=self._ctx.terminal_state, @@ -1060,7 +1061,7 @@ def _export_context(self) -> jtrans.TranslatedJaxprSDFG: def _handle_null_jaxpr( self, - jaxpr: jcore.ClosedJaxpr, + jaxpr: jax_core.ClosedJaxpr, ) -> Sequence[str]: """This function is called in case a `Jaxpr` with zero equations is encountered. @@ -1092,7 +1093,7 @@ def _handle_null_jaxpr( # Thus we have to introduce a some fake output name and explicitly copy the data around. # Once DaCe will inline the nested SDFG it will remove this intermediate copy. for jax_out_var in jaxpr.jaxpr.outvars: - jax_inp_name = jutil.get_jax_var_name( + jax_inp_name = util.get_jax_var_name( jax_out_var ) # Since output == input their names must be the same. assert self.map_jax_var_to_sdfg(jax_inp_name, allow_fail=True) diff --git a/src/jace/translator/sub_translators/__init__.py b/src/jace/translator/sub_translators/__init__.py index 7019076..88c239c 100644 --- a/src/jace/translator/sub_translators/__init__.py +++ b/src/jace/translator/sub_translators/__init__.py @@ -4,25 +4,24 @@ # All rights reserved. # # SPDX-License-Identifier: BSD-3-Clause - """Module collecting all built-in subtranslators.""" from __future__ import annotations from collections.abc import Sequence -from jace import translator as jtrans -from jace.translator.sub_translators.alu_translator import ALUTranslator +from .a_primitive_translator import PrimitiveTranslator # has to be the first import. +from .alu_translator import ALUTranslator # List of all subtranslators that ships with JaCe. -_KNOWN_SUBTRANSLATORS: list[type[jtrans.PrimitiveTranslator]] = [ +_KNOWN_SUBTRANSLATORS: list[type[PrimitiveTranslator]] = [ ALUTranslator, ] def add_subtranslator( - subtrans: type[jtrans.PrimitiveTranslator], + subtrans: type[PrimitiveTranslator], ) -> bool: """Add `subtrans` to the externally defined subtranslators. @@ -37,7 +36,7 @@ def add_subtranslator( return True -def _get_subtranslators_cls() -> Sequence[type[jtrans.PrimitiveTranslator]]: +def _get_subtranslators_cls() -> Sequence[type[PrimitiveTranslator]]: """Returns the list of all subtranslator known to JaCe. The translators are returned in FIFO order. @@ -48,4 +47,5 @@ def _get_subtranslators_cls() -> Sequence[type[jtrans.PrimitiveTranslator]]: __all__ = [ "ALUTranslator", "add_subtranslator", + "PrimitiveTranslator", ] diff --git a/src/jace/translator/primitive_translator.py b/src/jace/translator/sub_translators/a_primitive_translator.py similarity index 90% rename from src/jace/translator/primitive_translator.py rename to src/jace/translator/sub_translators/a_primitive_translator.py index 816b709..16d2d4b 100644 --- a/src/jace/translator/primitive_translator.py +++ b/src/jace/translator/sub_translators/a_primitive_translator.py @@ -4,6 +4,14 @@ # All rights reserved. # # SPDX-License-Identifier: BSD-3-Clause +"""Contains the interface for all primitive subtranslators. + +Note the name of this file is because it has to be the first that is imported in the `__init__.py` file. +If not, we would get a cyclic import error. +However, all attempts to prevent ruff from mindlessly (rule abiding) destroying this orders failed. +Thus the name was changed to enforce this. +If you have the solution, feel free to implement it. +""" from __future__ import annotations @@ -12,7 +20,7 @@ from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable import dace -from jax import core as jcore +from jax import core as jax_core if TYPE_CHECKING: @@ -21,7 +29,7 @@ @runtime_checkable class PrimitiveTranslator(Protocol): - """Interface for all Jax primitive subtranslators. + """Interface for all Jax primitive translators, also known as subtranslator. A translator for a primitive translates a single equation of a Jaxpr into its SDFG equivalent. A type that implements this interface must fulfil the following properties: @@ -66,7 +74,7 @@ def translate_jaxeqn( driver: JaxprTranslationDriver, in_var_names: Sequence[str | None], out_var_names: Sequence[str], - eqn: jcore.JaxprEqn, + eqn: jax_core.JaxprEqn, eqn_state: dace.SDFGState, ) -> dace.SDFGState | None: """Translates the Jax primitive into its SDFG equivalent. diff --git a/src/jace/translator/sub_translators/alu_translator.py b/src/jace/translator/sub_translators/alu_translator.py index b08e474..f397bb3 100644 --- a/src/jace/translator/sub_translators/alu_translator.py +++ b/src/jace/translator/sub_translators/alu_translator.py @@ -14,14 +14,13 @@ import dace import numpy as np -from jax import core as jcore +from jax import core as jax_core from typing_extensions import override -from jace import translator as jtranslator +from jace.translator import sub_translators -class ALUTranslator(jtranslator.PrimitiveTranslator): - # class ALUTranslator(PrimitiveTranslator): +class ALUTranslator(sub_translators.PrimitiveTranslator): """This translator handles all arithmetic and logical operations.""" __slots__ = () @@ -92,10 +91,10 @@ def primitive(self) -> Sequence[str]: @override def translate_jaxeqn( self, - driver: jtranslator.JaxprTranslationDriver, + driver: sub_translators.JaxprTranslationDriver, in_var_names: Sequence[str | None], out_var_names: Sequence[str], - eqn: jcore.JaxprEqn, + eqn: jax_core.JaxprEqn, eqn_state: dace.SDFGState, ) -> None: """Perform the translation. @@ -244,7 +243,7 @@ def translate_jaxeqn( def _writeTaskletCode( self, in_var_names: Sequence[str | None], - eqn: jcore.JaxprEqn, + eqn: jax_core.JaxprEqn, ) -> str: """This function generates the Tasklet code based on a primitive. @@ -284,7 +283,7 @@ def _writeTaskletCode( if in_var_name is not None: continue - jax_in_var: jcore.Literal = cast(jcore.Literal, eqn.invars[i]) + jax_in_var: jax_core.Literal = cast(jax_core.Literal, eqn.invars[i]) if jax_in_var.aval.shape == (): t_val = jax_in_var.val if isinstance(t_val, np.ndarray): diff --git a/src/jace/translator/translated_jaxpr_sdfg.py b/src/jace/translator/translated_jaxpr_sdfg.py index bc1ffab..3a3bb6b 100644 --- a/src/jace/translator/translated_jaxpr_sdfg.py +++ b/src/jace/translator/translated_jaxpr_sdfg.py @@ -12,9 +12,9 @@ from typing import Any import dace -from jax import core as jcore +from jax import core as jax_core -from jace import util as jutil +from jace import util @dataclass(init=True, repr=True, eq=False, frozen=False, kw_only=True, slots=True) @@ -37,7 +37,7 @@ class TranslatedJaxprSDFG: """ sdfg: dace.SDFG - jax_name_map: Mapping[jcore.Var | jutil.JaCeVar, str] + jax_name_map: Mapping[jax_core.Var | util.JaCeVar, str] start_state: dace.SDFGState | None = None terminal_state: dace.SDFGState | None = None inp_names: Sequence[str] | None = None diff --git a/src/jace/util/debug.py b/src/jace/util/debug.py index b3a84f3..27c57aa 100644 --- a/src/jace/util/debug.py +++ b/src/jace/util/debug.py @@ -13,18 +13,16 @@ from __future__ import annotations from collections.abc import Callable -from typing import TYPE_CHECKING, Any +from typing import Any import dace import jax - -if TYPE_CHECKING: - from jace import translator as jtrans +from jace import translator def run_jax_sdfg( - jsdfg: jtrans.TranslatedJaxprSDFG, + jsdfg: translator.TranslatedJaxprSDFG, *args: Any, ) -> tuple[Any, ...] | Any: """Calls the SDFG that is encapsulated with the supplied arguments. @@ -97,9 +95,7 @@ def _jace_run( *args: Forwarded to the tracing and final execution of the SDFG. **kwargs: Used to construct the driver. """ - from jace.translator import JaxprTranslationDriver - jaxpr = jax.make_jaxpr(fun)(*args) - driver = JaxprTranslationDriver(**kwargs) + driver = translator.JaxprTranslationDriver(**kwargs) jsdfg = driver.translate_jaxpr(jaxpr) return run_jax_sdfg(jsdfg, *args) diff --git a/src/jace/util/jax_helper.py b/src/jace/util/jax_helper.py index 5bd7596..80018fc 100644 --- a/src/jace/util/jax_helper.py +++ b/src/jace/util/jax_helper.py @@ -17,10 +17,10 @@ from collections.abc import Mapping from dataclasses import dataclass -from typing import Any +from typing import Any, overload import dace -import jax.core as jcore +import jax.core as jax_core import numpy as np from jace import util @@ -36,14 +36,13 @@ class JaCeVar: Notes: Main intention is to test functionality. - While for a Jax `Var` object the name is rather irrelevant, `JaCeVar` use their name. If the name of a `JaCeVar` is '_' it is considered a drop variable. If the name of a `JaCeVar` is empty, the automatic naming will consider it as a Jax variable. The definition of `__hash__` and `__eq__` is in accordance how Jax variable works. """ name: str - shape: tuple[int | dace.symbol | str, ...] | int | dace.symbol | str | tuple[()] + shape: tuple[int | dace.symbol | str, ...] | tuple[()] dtype: dace.typeclass def __hash__(self) -> int: @@ -59,7 +58,7 @@ def __post_init__(self) -> None: raise ValueError("The 'shape' member of a 'JaCeVar' must be a tuple.") -def get_jax_var_name(jax_var: jcore.Atom | JaCeVar | str) -> str: +def get_jax_var_name(jax_var: jax_core.Atom | JaCeVar | str) -> str: """Returns the name of the Jax variable as a string. Args: @@ -70,19 +69,19 @@ def get_jax_var_name(jax_var: jcore.Atom | JaCeVar | str) -> str: This function is subject for removal. """ match jax_var: - case jcore.DropVar(): + case jax_core.DropVar(): return "_" case JaCeVar(): # In case of an empty name consider the jace variable as a Jax variable. # This is mostly for testing. jax_name = f"jax{id(jax_var)}" if jax_var.name == "" else jax_var.name - case jcore.Var(): + case jax_core.Var(): # This stopped working after version 0.20.4, because of some changes in Jax # See `https://github.com/google/jax/pull/10573` for more information. # The following implementation will generate stable names, however, they will be decoupled # from output of the pretty printed Jaxpr jax_name = f"jax{jax_var.count}{jax_var.suffix}" - case jcore.Literal(): + case jax_core.Literal(): raise TypeError("Can not derive a name from a Jax Literal.") case str(): jax_name = jax_var @@ -97,14 +96,24 @@ def get_jax_var_name(jax_var: jcore.Atom | JaCeVar | str) -> str: return jax_name -def get_jax_var_shape(jax_var: jcore.Atom | JaCeVar) -> tuple[int, ...]: +@overload +def get_jax_var_shape(jax_var: JaCeVar) -> tuple[int | dace.symbol | str, ...] | tuple[()]: ... + + +@overload +def get_jax_var_shape(jax_var: jax_core.Atom) -> tuple[int, ...] | tuple[()]: ... + + +def get_jax_var_shape( + jax_var: jax_core.Atom | JaCeVar, +) -> tuple[int | dace.symbol | str, ...] | tuple[()]: """Returns the shape of a Jax variable. Args: jax_var: The variable to process """ match jax_var: - case jcore.Var() | jcore.Literal(): + case jax_core.Var() | jax_core.Literal(): return jax_var.aval.shape case JaCeVar(): return jax_var.shape @@ -112,10 +121,10 @@ def get_jax_var_shape(jax_var: jcore.Atom | JaCeVar) -> tuple[int, ...]: raise TypeError(f"'get_jax_var_shape()` is not implemented for '{type(jax_var)}'.") -def get_jax_var_dtype(jax_var: jcore.Atom | JaCeVar) -> dace.typeclass: +def get_jax_var_dtype(jax_var: jax_core.Atom | JaCeVar) -> dace.typeclass: """Returns the DaCe equivalent of `jax_var`s datatype.""" match jax_var: - case jcore.Var() | jcore.Literal(): + case jax_core.Var() | jax_core.Literal(): return translate_dtype(jax_var.aval.dtype) case JaCeVar(): return translate_dtype(jax_var.dtype) @@ -148,8 +157,8 @@ def translate_dtype(dtype: Any) -> dace.typeclass: def _propose_jax_name( - jax_var: jcore.Atom | JaCeVar, - jax_name_map: Mapping[jcore.Var | JaCeVar, Any] | None = None, + jax_var: jax_core.Atom | JaCeVar, + jax_name_map: Mapping[jax_core.Var | JaCeVar, Any] | None = None, ) -> str: """Proposes a variable name for `jax_var`. @@ -167,11 +176,9 @@ def _propose_jax_name( The naming of variables are only consistent with the inner most Jaxpr a variable is defined in. Dropped variables will always be named `'_'`. """ - from jace.util.traits import is_drop_var - - if is_drop_var(jax_var): + if util.traits.is_drop_var(jax_var): return "_" - if isinstance(jax_var, jcore.Literal): + if isinstance(jax_var, jax_core.Literal): raise TypeError(f"Can not propose a name for literal '{jax_var}'.") if jax_name_map is None: return get_jax_var_name(jax_var) @@ -180,7 +187,7 @@ def _propose_jax_name( raise RuntimeError( f"Can not propose a second name for '{jax_var}', it already known as '{jax_name_map[jax_var]}'." ) - if isinstance(jax_var, jcore.Var): + if isinstance(jax_var, jax_core.Var): pass elif isinstance(jax_var, JaCeVar): # If the name of the JaCe variable is empty, then use the name proposing diff --git a/src/jace/util/traits.py b/src/jace/util/traits.py index 247c999..1e063c8 100644 --- a/src/jace/util/traits.py +++ b/src/jace/util/traits.py @@ -12,9 +12,9 @@ from collections.abc import Iterable from typing import Any, TypeGuard -from jax import core as jcore +from jax import core as jax_core -from jace import util as jutil +from jace import util class NonStringIterable(Iterable): ... @@ -24,11 +24,11 @@ def is_non_string_iterable(val: Any) -> TypeGuard[NonStringIterable]: return isinstance(val, Iterable) and not isinstance(val, str) -def is_drop_var(jax_var: jcore.Atom | jutil.JaCeVar) -> bool: +def is_drop_var(jax_var: jax_core.Atom | util.JaCeVar) -> bool: """Tests if `jax_var` is a drop variable.""" - if isinstance(jax_var, jcore.DropVar): + if isinstance(jax_var, jax_core.DropVar): return True - if isinstance(jax_var, jutil.JaCeVar): + if isinstance(jax_var, util.JaCeVar): return jax_var.name == "_" return False diff --git a/src/jace/util/util.py b/src/jace/util/util.py index 3943743..96bfa20 100644 --- a/src/jace/util/util.py +++ b/src/jace/util/util.py @@ -10,6 +10,8 @@ from collections.abc import Iterable from typing import TypeVar, cast, overload +from jace.util import traits + _T = TypeVar("_T") @@ -27,8 +29,6 @@ def as_sequence(value: _T) -> Iterable[_T]: ... def as_sequence(value: _T | Iterable[_T]) -> Iterable[_T]: - from jace.util.traits import is_non_string_iterable - - if is_non_string_iterable(value): + if traits.is_non_string_iterable(value): return value return cast(Iterable[_T], [value])