Merge pull request #306 from superbobry:maint-2

PiperOrigin-RevId: 688121597

jax_triton/experimental/fusion/jaxpr_rewriter.py

@@ -120,7 +120,7 @@ def shape(self):
 
   def match(self, expr, bindings, succeed):
     if not isinstance(expr, Literal):
-      return []
+      return []  # noqa: B901
     yield from matcher.matcher((self.value, self.dtype))((expr.value,
       expr.dtype), bindings, succeed)
 
@@ -138,7 +138,7 @@ class Part(Node):
 
   def match(self, expr, bindings, succeed):
     if not isinstance(expr, Part):
-      return []
+      return []  # noqa: B901
     yield from matcher.matcher((self.index, self.shape, self.dtype, self.parent))((
       expr.index, expr.shape, expr.dtype, expr.parent), bindings, succeed)
 

jax_triton/triton_lib.py

@@ -51,7 +51,6 @@
   import triton.language as tl
   from triton.runtime import autotuner
   import triton._C.libtriton as _triton
-  from triton._C.libtriton import ir as tl_ir
   import triton.backends.nvidia.compiler as cb
 
   CAN_USE_TRITON = True
@@ -375,7 +374,7 @@ def get_or_create_triton_kernel(
     args_for_specialization_attr[i] = v
   specialization_attr = fn._get_config(*args_for_specialization_attr)  # pylint: disable=protected-access
 
-  constants = {k: v for k, v in metaparams.items()}
+  constants = dict(metaparams)
   constants.update({k: None for _, k, v in scalar_args if v is None})
   constants.update({fn.arg_names[i]: 1 for i in specialization_attr.equal_to_1})
 

pyproject.toml

@@ -40,10 +40,34 @@ target-version = "py310"
 
 [tool.ruff.lint]
 ignore = [
-    # Do not assign a `lambda` expression, use a `def`
-    "E731",
-    # Module level import not at top of file
-    "E402",
-    # Ambiguous variable name
-    "E741",
+    # Unnecessary collection call
+    "C408",
+    # Unnecessary map usage
+    "C417",
+    # Unnecessary dict comprehension for iterable
+    "C420",
+    # Object names too complex
+    "C901",
+    # Local variable is assigned to but never used
+    "F841",
+    # Raise with from clause inside except block
+    "B904",
+    # Zip without explicit strict parameter
+    "B905",
+]
+select = [
+    "B9",
+    "C",
+    "F",
+    "W",
+    "YTT",
+    "ASYNC",
+    "E101",
+    "E112",
+    "E113",
+    "E115",
+    "E117",
+    "E225",
+    "E227",
+    "E228",
 ]