Removed inspect.signature() call from jaxlib.triton.dialect.ScanOp

PiperOrigin-RevId: 614772594

Author: superbobry

jax/_src/pallas/triton/lowering.py

@@ -492,7 +492,7 @@ def _associative_scan_lowering(body, ctx: LoweringRuleContext, args, axes):
   if consts:
     raise NotImplementedError("Associative scan with constants not supported.")
   element_types = [_element_type(arg.type) for arg in flat_args]
-  scan_op = tt_dialect.ScanOp(flat_args, axis, reverse=False)
+  scan_op = tt_dialect.ScanOp(flat_args, axis)
   param_types = element_types * 2
   entry = scan_op.regions[0].blocks.append(*param_types)
   with ir.InsertionPoint.at_block_begin(entry):

jaxlib/triton/dialect.py

@@ -19,7 +19,6 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-import inspect
 
 from jaxlib.mlir._mlir_libs._triton_ext import (
     PointerType,
@@ -56,18 +55,13 @@ def __init__(
       self,
       operands: Sequence[ir.Value],
       axis: int,
-      reverse: ir.Attribute,
+      reverse: bool = False,
       *,
       loc: ir.Location | None = None,
       ip: ir.InsertionPoint | None = None,
   ):
     return_types = [op.type for op in operands]
-    # OSS might have an old version of Triton, whose ScanOp doesn't have a
-    # reverse parameter.
-    if "reverse" in inspect.signature(super().__init__).parameters:
-      super().__init__(return_types, operands, axis, reverse, loc=loc, ip=ip)
-    else:
-      super().__init__(return_types, operands, axis, loc=loc, ip=ip)  # pylint: disable=no-value-for-parameter
+    super().__init__(return_types, operands, axis, reverse, loc=loc, ip=ip)
 
 
 # TODO(slebedev): Consider overriding instead.