diff --git a/mlir/python/mlir/dialects/scf.py b/mlir/python/mlir/dialects/scf.py index 9e22df3dd50a9..27e8c6f5dd05f 100644 --- a/mlir/python/mlir/dialects/scf.py +++ b/mlir/python/mlir/dialects/scf.py @@ -193,11 +193,11 @@ def block(self) -> Block: class IfOp(IfOp): """Specialization for the SCF if op class.""" - def __init__(self, cond, results_=None, *, hasElse=False, loc=None, ip=None): + def __init__(self, cond, results_=None, *, has_else=False, loc=None, ip=None): """Creates an SCF `if` operation. - `cond` is a MLIR value of 'i1' type to determine which regions of code will be executed. - - `hasElse` determines whether the if operation has the else branch. + - `has_else` determines whether the if operation has the else branch. """ if results_ is None: results_ = [] @@ -207,17 +207,19 @@ def __init__(self, cond, results_=None, *, hasElse=False, loc=None, ip=None): results.extend(results_) super().__init__(results, cond, loc=loc, ip=ip) self.regions[0].blocks.append(*[]) - if hasElse: + if has_else: self.regions[1].blocks.append(*[]) @property - def then_block(self): + def then_block(self) -> Block: """Returns the then block of the if operation.""" return self.regions[0].blocks[0] @property - def else_block(self): + def else_block(self) -> Optional[Block]: """Returns the else block of the if operation.""" + if len(self.regions[1].blocks) == 0: + return None return self.regions[1].blocks[0] diff --git a/mlir/test/python/dialects/scf.py b/mlir/test/python/dialects/scf.py index 0c0c9b986562b..e53a299d1f869 100644 --- a/mlir/test/python/dialects/scf.py +++ b/mlir/test/python/dialects/scf.py @@ -335,7 +335,7 @@ def testIfWithElse(): @func.FuncOp.from_py_func(bool) def simple_if_else(cond): - if_op = scf.IfOp(cond, [i32, i32], hasElse=True) + if_op = scf.IfOp(cond, [i32, i32], has_else=True) with InsertionPoint(if_op.then_block): x_true = arith.ConstantOp(i32, 0) y_true = arith.ConstantOp(i32, 1)