Merge branch 'new_translators' into new_test_suite

src/jace/translator/primitive_translators/convert_element_type_translator.py

@@ -32,7 +32,7 @@ class ConvertElementTypeTranslator(mapped_base.MappedOperationTranslatorBase):
     will perform the type conversion operation.
 
     Note:
-        The type to cast to id inferred from the output variable and the `new_dtype`
+        The type to cast to is inferred from the output variable and the `new_dtype`
         parameter of the equation is ignored.
     """
 

src/jace/translator/primitive_translators/gather_translator.py

@@ -63,7 +63,7 @@ def gather_translator(  # noqa: PLR0914 [too-many-locals]  # Can not reduce any
     if eqn.params["mode"] != jax_lax.GatherScatterMode.PROMISE_IN_BOUNDS:
         raise NotImplementedError(f"The mode {eqn.params['mode']} is not implemented.")
 
-    # This is the size of the slice window that is copied. Its length equal the rank
+    # This is the size of the slice window that is copied. Its length is the rank
     #  of the source array, dimensions that are excluded from copying are listed
     #  in `collapsed_slice_dims`.
     slice_sizes: Sequence[int] = eqn.params["slice_sizes"]

src/jace/translator/primitive_translators/pjit_translator.py

@@ -57,9 +57,7 @@ def pjit_translator(
     nested_jaxpr: jax_core.ClosedJaxpr = eqn.params["jaxpr"]
     in_shardings = eqn.params["in_shardings"]
     out_shardings = eqn.params["out_shardings"]
-    _ = eqn.params["donated_invars"]  # Always ignored
-    _ = eqn.params["keep_unused"]
-    _ = eqn.params["inline"]
+    # "donated_invars", "keep_unused", "inline" parameters are just ignored
 
     if not all(in_sharding is jax_sharding.UNSPECIFIED for in_sharding in in_shardings):
         raise NotImplementedError("Currently 'pjit' does not support sharding in its input.")