WIP: reformat MLIR python with yapf

mlir/python/mlir/_mlir_libs/__init__.py

def get_include_dirs() -> Sequence[str]:
   Depending on how the package was build, development C libraries may or may
   not be present.
   """
   return [os.path.join(_this_dir, "include")]
 
 
 # Perform Python level site initialization. This involves:
 #   1. Attempting to load initializer modules, specific to the distribution.
-#   2. Defining the concrete mlir.ir.Context that does site specific 
+#   2. Defining the concrete mlir.ir.Context that does site specific
 #      initialization.
 #
 # Aside from just being far more convenient to do this at the Python level,
 # it is actually quite hard/impossible to have such __init__ hooks, given
 # the pybind memory model (i.e. there is not a Python reference to the object
 # in the scope of the base class __init__).
 #
 # For #1, we:
-#   a. Probe for modules named '_mlirRegisterEverything' and 
-#     '_site_initialize_{i}', where 'i' is a number starting at zero and 
+#   a. Probe for modules named '_mlirRegisterEverything' and
+#     '_site_initialize_{i}', where 'i' is a number starting at zero and
 #     proceeding so long as a module with the name is found.
 #   b. If the module has a 'register_dialects' attribute, it will be called
 #     immediately with a DialectRegistry to populate.
 #   c. If the module has a 'context_init_hook', it will be added to a list
-#     of callbacks that are invoked as the last step of Context 
+#     of callbacks that are invoked as the last step of Context
 #     initialization (and passed the Context under construction).
 #
 # This facility allows downstreams to customize Context creation to their
 # needs.
 def _site_initialize():
   import importlib
   import itertools
   import logging
   from ._mlir import ir
   logger = logging.getLogger(__name__)
   registry = ir.DialectRegistry()
   post_init_hooks = []
 
   def process_initializer_module(module_name):
     try:
       m = importlib.import_module(f".{module_name}", __name__)
     except ModuleNotFoundError:
       return False
     except ImportError:
-      message = (f"Error importing mlir initializer {module_name}. This may "
+      message = (
+          f"Error importing mlir initializer {module_name}. This may "
           "happen in unclean incremental builds but is likely a real bug if "
           "encountered otherwise and the MLIR Python API may not function.")
       logger.warning(message, exc_info=True)
 
     logger.debug("Initializing MLIR with module: %s", module_name)
     if hasattr(m, "register_dialects"):
       logger.debug("Registering dialects from initializer %r", m)
       m.register_dialects(registry)
     if hasattr(m, "context_init_hook"):
       logger.debug("Adding context init hook from %r", m)
       post_init_hooks.append(m.context_init_hook)
     return True
 
-
   # If _mlirRegisterEverything is built, then include it as an initializer
   # module.
   process_initializer_module("_mlirRegisterEverything")
 
   # Load all _site_initialize_{i} modules, where 'i' is a number starting
   # at 0.
   for i in itertools.count():
     module_name = f"_site_initialize_{i}"
     if not process_initializer_module(module_name):
       break
 
   class Context(ir._BaseContext):
+
     def __init__(self, *args, **kwargs):
       super().__init__(*args, **kwargs)
       self.append_dialect_registry(registry)
       for hook in post_init_hooks:
         hook(self)
       # TODO: There is some debate about whether we should eagerly load
       # all dialects. It is being done here in order to preserve existing
       # behavior. See: https://github.com/llvm/llvm-project/issues/56037
       self.load_all_available_dialects()
+
   ir.Context = Context
 
   class MLIRError(Exception):
     """
     An exception with diagnostic information. Has the following fields:
       message: str
       error_diagnostics: List[ir.DiagnosticInfo]
     """
+
     def __init__(self, message, error_diagnostics):
       self.message = message
       self.error_diagnostics = error_diagnostics
       super().__init__(message, error_diagnostics)
 
     def __str__(self):
       s = self.message
       if self.error_diagnostics:
         s += ':'
       for diag in self.error_diagnostics:
-        s += "\nerror: "  + str(diag.location)[4:-1] + ": " + diag.message.replace('\n', '\n  ')
+        s += "\nerror: " + str(
+            diag.location)[4:-1] + ": " + diag.message.replace('\n', '\n  ')
         for note in diag.notes:
-          s += "\n note: "  + str(note.location)[4:-1] + ": " + note.message.replace('\n', '\n  ')
+          s += "\n note: " + str(
+              note.location)[4:-1] + ": " + note.message.replace('\n', '\n  ')
       return s
+
   ir.MLIRError = MLIRError
 
 
 _site_initialize()

mlir/python/mlir/dialects/_arith_ops_ext.py

def __init__(self,
     elif isinstance(value, float):
       super().__init__(FloatAttr.get(result, value), loc=loc, ip=ip)
     else:
       super().__init__(value, loc=loc, ip=ip)
 
   @classmethod
   def create_index(cls, value: int, *, loc=None, ip=None):
     """Create an index-typed constant."""
-    return cls(
-        IndexType.get(context=_get_default_loc_context(loc)),
+    return cls(IndexType.get(context=_get_default_loc_context(loc)),
                value,
                loc=loc,
                ip=ip)
 
   @property
   def type(self):
     return self.results[0].type
 
   @property
   def value(self):
     return Attribute(self.operation.attributes["value"])

mlir/python/mlir/dialects/_bufferization_ops_ext.py

def __init__(self,
                *,
                loc=None,
                ip=None):
     """Constructs an `alloc_tensor` with static and/or dynamic sizes."""
     context = get_default_loc_context(loc)
     attributes = {}
     if escape:
       attributes["escape"] = escape
-    op = self.build_generic(
-        results=[tensor_type],
+    op = self.build_generic(results=[tensor_type],
                             operands=[dynamic_sizes, copy, size_hint],
                             attributes=attributes,
                             loc=loc,
                             ip=ip)
     OpView.__init__(self, op)

mlir/python/mlir/dialects/_builtin_ops_ext.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 try:
   from ..ir import *
 except ImportError as e:
   raise RuntimeError("Error loading imports from extension module") from e
 
+
 class ModuleOp:
   """Specialization for the module op class."""
 
   def __init__(self, *, loc=None, ip=None):
     super().__init__(self.build_generic(results=[], operands=[], loc=loc,
                                         ip=ip))
     body = self.regions[0].blocks.append()
 
   @property
   def body(self):
     return self.regions[0].blocks[0]

mlir/python/mlir/dialects/_func_ops_ext.py

try:
 
   from typing import Any, List, Optional, Sequence, Union
 except ImportError as e:
   raise RuntimeError("Error loading imports from extension module") from e
 
 ARGUMENT_ATTRIBUTE_NAME = "arg_attrs"
 RESULT_ATTRIBUTE_NAME = "res_attrs"
 
+
 class ConstantOp:
   """Specialization for the constant op class."""
 
   def __init__(self, result: Type, value: Attribute, *, loc=None, ip=None):
     super().__init__(result, value, loc=loc, ip=ip)
 
   @property
   def type(self):
def __init__(self,
 
     # If the type is passed as a tuple, build a FunctionType on the fly.
     if isinstance(type, tuple):
       type = FunctionType.get(inputs=type[0], results=type[1])
 
     type = TypeAttr.get(type)
     sym_visibility = StringAttr.get(
         str(visibility)) if visibility is not None else None
-    super().__init__(sym_name, type, sym_visibility=sym_visibility, loc=loc, ip=ip)
+    super().__init__(sym_name,
+                     type,
+                     sym_visibility=sym_visibility,
+                     loc=loc,
+                     ip=ip)
     if body_builder:
       entry_block = self.add_entry_block()
       with InsertionPoint(entry_block):
         body_builder(self)
 
   @property
   def is_external(self):
     return len(self.regions[0].blocks) == 0
def decorator(f):
 
       wrapped = emit_call_op
       wrapped.__name__ = f.__name__
       wrapped.func_op = func_op
       return wrapped
 
     return decorator
 
+
 class CallOp:
   """Specialization for the call op class."""
 
   def __init__(self,
                calleeOrResults: Union[FuncOp, List[Type]],
                argumentsOrCallee: Union[List, FlatSymbolRefAttr, str],
                arguments: Optional[List] = None,
                *,
if isinstance(calleeOrResults, FuncOp):
         raise ValueError(
             "when constructing a call to a function, expected " +
             "the second argument to be a list of call arguments, " +
             f"got {type(argumentsOrCallee)}")
       if arguments is not None:
         raise ValueError("unexpected third argument when constructing a call" +
                          "to a function")
 
-      super().__init__(
-          calleeOrResults.type.results,
+      super().__init__(calleeOrResults.type.results,
                        FlatSymbolRefAttr.get(
                            calleeOrResults.name.value,
                            context=_get_default_loc_context(loc)),
                        argumentsOrCallee,
                        loc=loc,
                        ip=ip)
       return
 
     if isinstance(argumentsOrCallee, list):
       raise ValueError("when constructing a call to a function by name, " +
                        "expected the second argument to be a string or a " +
                        f"FlatSymbolRefAttr, got {type(argumentsOrCallee)}")
 
     if isinstance(argumentsOrCallee, FlatSymbolRefAttr):
-      super().__init__(
-          calleeOrResults, argumentsOrCallee, arguments, loc=loc, ip=ip)
+      super().__init__(calleeOrResults,
+                       argumentsOrCallee,
+                       arguments,
+                       loc=loc,
+                       ip=ip)
     elif isinstance(argumentsOrCallee, str):
-      super().__init__(
-          calleeOrResults,
+      super().__init__(calleeOrResults,
                        FlatSymbolRefAttr.get(
-              argumentsOrCallee, context=_get_default_loc_context(loc)),
+                           argumentsOrCallee,
+                           context=_get_default_loc_context(loc)),
                        arguments,
                        loc=loc,
                        ip=ip)

mlir/python/mlir/dialects/_linalg_ops_ext.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 try:
   from typing import Optional, Sequence, Union
   from ..ir import *
   from ._ods_common import get_default_loc_context
   from .._mlir_libs._mlirDialectsLinalg import fill_builtin_region
 except ImportError as e:
   raise RuntimeError("Error loading imports from extension module") from e
 
 from ._ods_common import get_op_result_or_value as _get_op_result_or_value
 
+
 def isa(cls: Type, ty: Type):
   try:
     cls(ty)
     return True
   except ValueError:
     return False
 
 

mlir/python/mlir/dialects/_ml_program_ops_ext.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 try:
   from typing import Union
   from ..ir import *
   from ._ods_common import get_default_loc_context as _get_default_loc_context
 except ImportError as e:
   raise RuntimeError("Error loading imports from extension module") from e
 
 from ._ml_program_ops_gen import *
 
-
 ARGUMENT_ATTRIBUTE_NAME = "arg_attrs"
 RESULT_ATTRIBUTE_NAME = "res_attrs"
 
 
 class FuncOp:
   """Specialization for the func op class."""
 
   def __init__(self,
def __init__(self,
 
     # If the type is passed as a tuple, build a FunctionType on the fly.
     if isinstance(type, tuple):
       type = FunctionType.get(inputs=type[0], results=type[1])
 
     type = TypeAttr.get(type)
     sym_visibility = StringAttr.get(
         str(visibility)) if visibility is not None else None
-    super().__init__(sym_name, type, sym_visibility=sym_visibility, loc=loc, ip=ip)
+    super().__init__(sym_name,
+                     type,
+                     sym_visibility=sym_visibility,
+                     loc=loc,
+                     ip=ip)
     if body_builder:
       entry_block = self.add_entry_block()
       with InsertionPoint(entry_block):
         body_builder(self)
 
   @property
   def is_external(self):
     return len(self.regions[0].blocks) == 0

mlir/python/mlir/dialects/_ods_common.py

def get_default_loc_context(location=None):
   """
   if location is None:
     # Location.current raises ValueError if there is no current location.
     return _cext.ir.Location.current.context
   return location.context
 
 
 def get_op_result_or_value(
-    arg: _Union[_cext.ir.OpView, _cext.ir.Operation, _cext.ir.Value, _cext.ir.OpResultList]
+    arg: _Union[_cext.ir.OpView, _cext.ir.Operation, _cext.ir.Value,
+                _cext.ir.OpResultList]
 ) -> _cext.ir.Value:
   """Returns the given value or the single result of the given op.
 
   This is useful to implement op constructors so that they can take other ops as
   arguments instead of requiring the caller to extract results for every op.
   Raises ValueError if provided with an op that doesn't have a single result.
   """
   if isinstance(arg, _cext.ir.OpView):
     return arg.operation.result
   elif isinstance(arg, _cext.ir.Operation):
     return arg.result
   elif isinstance(arg, _cext.ir.OpResultList):
     return arg[0]
   else:
     assert isinstance(arg, _cext.ir.Value)
     return arg
 
 
 def get_op_results_or_values(
     arg: _Union[_cext.ir.OpView, _cext.ir.Operation,
-                _Sequence[_Union[_cext.ir.OpView, _cext.ir.Operation, _cext.ir.Value]]]
+                _Sequence[_Union[_cext.ir.OpView, _cext.ir.Operation,
+                                 _cext.ir.Value]]]
 ) -> _Union[_Sequence[_cext.ir.Value], _cext.ir.OpResultList]:
   """Returns the given sequence of values or the results of the given op.
 
   This is useful to implement op constructors so that they can take other ops as
   lists of arguments instead of requiring the caller to extract results for
   every op.
   """
   if isinstance(arg, _cext.ir.OpView):
     return arg.operation.results
   elif isinstance(arg, _cext.ir.Operation):
     return arg.results
   else:
     return [get_op_result_or_value(element) for element in arg]

mlir/python/mlir/dialects/_scf_ops_ext.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 try:
   from ..ir import *
 except ImportError as e:
   raise RuntimeError("Error loading imports from extension module") from e
 
 from typing import Any, Optional, Sequence, Union
 from ._ods_common import get_op_result_or_value as _get_op_result_or_value, get_op_results_or_values as _get_op_results_or_values
 
+
 class ForOp:
   """Specialization for the SCF for op class."""
 
   def __init__(self,
                lower_bound,
                upper_bound,
                step,
                iter_args: Optional[Union[Operation, OpView,
- `iter_args` is a list of additional loop-carried arguments or an operation
       producing them as results.
     """
     if iter_args is None:
       iter_args = []
     iter_args = _get_op_results_or_values(iter_args)
 
     results = [arg.type for arg in iter_args]
     super().__init__(
-        self.build_generic(
-            regions=1,
+        self.build_generic(regions=1,
                            results=results,
                            operands=[
                                _get_op_result_or_value(o)
                                for o in [lower_bound, upper_bound, step]
                            ] + list(iter_args),
                            loc=loc,
                            ip=ip))
     self.regions[0].blocks.append(IndexType.get(), *results)
 
   @property
   def body(self):
     """Returns the body (block) of the loop."""
     return self.regions[0].blocks[0]
 
   @property
   def induction_variable(self):
     """Returns the induction variable of the loop."""
     return self.body.arguments[0]
 
   @property
   def inner_iter_args(self):
     """Returns the loop-carried arguments usable within the loop.
 
     To obtain the loop-carried operands, use `iter_args`.
     """
     return self.body.arguments[1:]
 
 
 class IfOp:
   """Specialization for the SCF if op class."""
 
-  def __init__(self,
-               cond,
-               results_=[],
-               *,
-               hasElse=False,
-               loc=None,
-               ip=None):
+  def __init__(self, cond, results_=[], *, hasElse=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.
     """
     operands = []
     operands.append(cond)
     results = []
     results.extend(results_)
     super().__init__(
-        self.build_generic(
-            regions=2,
+        self.build_generic(regions=2,
                            results=results,
                            operands=operands,
                            loc=loc,
                            ip=ip))
     self.regions[0].blocks.append(*[])
     if hasElse:
-        self.regions[1].blocks.append(*[])
 
   @property
   def then_block(self):
     """Returns the then block of the if operation."""
     return self.regions[0].blocks[0]
 
   @property
   def else_block(self):
     """Returns the else block of the if operation."""
     return self.regions[1].blocks[0]

mlir/python/mlir/dialects/_tensor_ops_ext.py

def __init__(self,
     static_sizes = []
     for s in sizes:
       if isinstance(s, int):
         static_sizes.append(s)
       else:
         static_sizes.append(ShapedType.get_dynamic_size())
         dynamic_sizes.append(s)
     result_type = RankedTensorType.get(static_sizes, element_type)
-    op = self.build_generic(
-        results=[result_type],
+    op = self.build_generic(results=[result_type],
                             operands=dynamic_sizes,
                             attributes={},
                             loc=loc,
                             ip=ip)
     OpView.__init__(self, op)

mlir/python/mlir/dialects/_transform_pdl_extension_ops_ext.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 try:
   from ..ir import *
   from ._ods_common import (
       get_op_result_or_value as _get_op_result_or_value,
       get_op_results_or_values as _get_op_results_or_values,
   )
 except ImportError as e:
   raise RuntimeError("Error loading imports from extension module") from e
 
 from typing import Union
 
+
 class PDLMatchOp:
 
   def __init__(
       self,
       result_type: Type,
       target: Union[Operation, Value],
       pattern_name: Union[Attribute, str],
       *,

mlir/python/mlir/dialects/linalg/opdsl/lang/comprehension.py

def __init__(self,
                output: bool = False):
     if index_dims and len(shape) != len(index_dims):
       raise ValueError(f"Expected the shape rank {len(shape)} to match the "
                        f"number of index_dims {len(index_dims)}")
     if index_dims and any(not isinstance(dim, DimDef) for dim in index_dims):
       raise ValueError(f"TensorDef requires index dims of type DimDef but "
                        f"got {index_dims}")
     kind = OperandKind.OUTPUT_TENSOR if output else OperandKind.INPUT_TENSOR
-    self.operand_def = OperandDef(
-        kind, type_var=type_var, size_exprs=shape, index_dims=index_dims)
+    self.operand_def = OperandDef(kind,
+                                  type_var=type_var,
+                                  size_exprs=shape,
+                                  index_dims=index_dims)
 
   def __getitem__(self, dims: Sequence[AffineExprDef]) -> TensorUse:
     assert self.operand_def.owner, "TensorDef is not registered with an op"
-    state = AffineBuildState(
-        global_state=self.operand_def.owner._affine_state,
+    state = AffineBuildState(global_state=self.operand_def.owner._affine_state,
                              allow_new_symbols=False)
     if not isinstance(dims, tuple):
       dims = (dims,)  # Handle single subscript case.
     # Special case: (None) is a 0d-scalar use.
     if dims == (None,):
       dims = ()
 
     exprs = []
     for expr_def in dims:
if any(not isinstance(size, SymbolDef) for size in sizes):
       raise ValueError(f"IndexAttrDef requires sizes of type SymbolDef "
                        f"but got {sizes}")
     if any(not isinstance(default_val, int) for default_val in default):
       raise ValueError(f"IndexAttrDef requires default values of type int "
                        f"but got {default}")
     if len(sizes) != len(default):
       raise ValueError(f"IndexAttrDef expects {len(sizes)} default values "
                        f"but got {len(default)}")
-    self.operand_def = OperandDef(
-        OperandKind.INDEX_ATTR, size_exprs=sizes, default_indices=default)
+    self.operand_def = OperandDef(OperandKind.INDEX_ATTR,
+                                  size_exprs=sizes,
+                                  default_indices=default)
 
 
 class UnaryFnAttrDef:
   """Unary function attribute definition.
 
   Unary function attributes provide a way to make the arithmetic computation
   parametrizable. Every attribute specifies a default unary function
   that may be overwritten at operation instantiation time.
   """
 
   def __init__(self, default: "UnaryFnType"):
     if not isinstance(default, UnaryFnType):
       raise ValueError(f"UnaryFnAttrDef requires default of type UnaryFnType "
                        f"but got {default}")
-    self.operand_def = OperandDef(
-        OperandKind.UNARY_FN_ATTR, default_fn=default.fn_name)
+    self.operand_def = OperandDef(OperandKind.UNARY_FN_ATTR,
+                                  default_fn=default.fn_name)
 
   def __call__(self, arg: TensorExpression) -> TensorFn:
     return TensorFn(FunctionKind.UNARY, None, self.operand_def, None, [arg])
 
 
 class BinaryFnAttrDef:
   """Binary function attribute definition.
 
   Binary function attributes provide a way to make the arithmetic computation
   parametrizable. Every attribute specifies a default binary function
   that may be overwritten at operation instantiation time.
   """
 
   def __init__(self, default: "BinaryFnType"):
     if not isinstance(default, BinaryFnType):
       raise ValueError(f"BinaryFnAttrDef requires default of type BinaryFnType "
                        f"but got {default}")
-    self.operand_def = OperandDef(
-        OperandKind.BINARY_FN_ATTR, default_fn=default.fn_name)
+    self.operand_def = OperandDef(OperandKind.BINARY_FN_ATTR,
+                                  default_fn=default.fn_name)
 
   def __call__(self, arg0: TensorExpression,
                arg1: TensorExpression) -> TensorFn:
     return TensorFn(FunctionKind.BINARY, None, self.operand_def, None,
                     [arg0, arg1])
 
   def __getitem__(self, reduce_dims: Tuple[DimDef]) -> ReduceFnUse:
     return ReduceFnUse(None, self, *reduce_dims)
 
 
 class TypeFnAttrDef:
   """Type conversion function attribute definition.
 
   Type conversion function attributes provide a way to make type conversions
   parameterizable. Every attribute specifies a default type conversion function
   that may be overwritten at operation instantiation time.
   """
 
   def __init__(self, default: "TypeFnType"):
     if not isinstance(default, TypeFnType):
       raise ValueError(f"TypeFnAttrDef requires default of type TypeFnType "
                        f"but got {default}")
-    self.operand_def = OperandDef(
-        OperandKind.TYPE_FN_ATTR, default_fn=default.fn_name)
+    self.operand_def = OperandDef(OperandKind.TYPE_FN_ATTR,
+                                  default_fn=default.fn_name)
 
   def __call__(self, type_var: TypeVar, arg: TensorExpression) -> TensorFn:
     return TensorFn(FunctionKind.TYPE, None, self.operand_def, type_var, [arg])
 
 
 ###############################################################################
 # Operation definition.
 ###############################################################################

 
 class LinalgOpDef:
   """Definition of a linalg op."""
 
   def __init__(self,
                name: str,
                cpp_class_name: Optional[str] = None,
                doc: Optional[str] = None):
-    self.metadata = OpMetadataDef(
-        name=name, cpp_class_name=cpp_class_name, doc=doc)
+    self.metadata = OpMetadataDef(name=name,
+                                  cpp_class_name=cpp_class_name,
+                                  doc=doc)
     self.registered_operands = dict()  # type: Dict[str, OperandDef]
     self.domain = list()  # type: List[DimDef]
     self.comprehensions = list()  # type: List[Comprehension]
     self._affine_state = AffineBuildState()
 
   def add_operand(self, name: str, operand: OperandDef):
     """Registers an operand."""
     if name in self.registered_operands:

mlir/python/mlir/dialects/linalg/opdsl/lang/config.py

def __init__(self,
     if (len(domain) != len(collected_dim_uses) or
         any(dim not in collected_dim_uses for dim in domain)):
       raise ValueError(f"Expected the annotated domain dimensions {domain} to "
                        f"match the set of dimension used by the tensor "
                        f"comprehension {collected_dim_uses}")
 
     # Instantiate the dimensions in the given order.
     with self.context:
-      local_state = AffineBuildState(
-          global_state=self.affine_state, allow_new_symbols=False)
+      local_state = AffineBuildState(global_state=self.affine_state,
+                                     allow_new_symbols=False)
       for dim in domain:
         dim.build(state=local_state)
 
     # Collect all attribute definitions.
     collected_attr_defs = list()
     for operand in registered_operands:
       if operand.is_attribute():
         collected_attr_defs.append(operand)
def __init__(self,
     # Generate the scalar assignments (used to build a body).
     self.assignments = [
         ScalarAssign(write_use.tensor_name, read_expr.to_scalar_expression())
         for write_use, read_expr in self.writes
     ]
 
   @property
   def ordered_operands(self) -> Sequence[OperandDefConfig]:
-    return sorted(
-        self.operands.values(),
+    return sorted(self.operands.values(),
                   key=lambda operand: operand.operand_def.registered_index)
 
   @property
   def ordered_dims(self) -> Sequence[Tuple[str, int]]:
     """Gets the ordered list of dim bindings (symbolic name, position).
 
     TODO: The original parser relies on parse ordering to arrive at the
     iterator types, but that ordering is not defined on the Python side, so
     this may be ambiguous.
class LinalgStructuredOpConfig(YAMLObject):
   def add_operand(self, operand_def: OperandDef):
     if operand_def in self.operands:
       return
     if not (operand_def.is_tensor() or
             operand_def.kind == OperandKind.INDEX_ATTR):
       self.operands[operand_def] = OperandDefConfig(operand_def)
       return
     with self.context:
-      local_state = AffineBuildState(
-          global_state=self.affine_state, allow_new_dims=False)
+      local_state = AffineBuildState(global_state=self.affine_state,
+                                     allow_new_dims=False)
       exprs = []
       for expr in operand_def.size_exprs:
         exprs.append(expr.build(state=local_state))
       assert local_state.local_dim_count == 0
-      affine_map = _ir.AffineMap.get(
-          dim_count=0, symbol_count=local_state.symbol_count, exprs=exprs)
+      affine_map = _ir.AffineMap.get(dim_count=0,
+                                     symbol_count=local_state.symbol_count,
+                                     exprs=exprs)
       if operand_def.kind == OperandKind.INDEX_ATTR:
-        self.operands[operand_def] = OperandDefConfig(
-            operand_def, index_attr_map=affine_map)
+        self.operands[operand_def] = OperandDefConfig(operand_def,
+                                                      index_attr_map=affine_map)
       else:
-        self.operands[operand_def] = OperandDefConfig(
-            operand_def, shape_map=affine_map)
+        self.operands[operand_def] = OperandDefConfig(operand_def,
+                                                      shape_map=affine_map)
 
   def add_indexed_operand(self, operand_def: OperandDef):
     with self.context:
-      local_state = AffineBuildState(
-          global_state=self.affine_state, allow_new_symbols=False)
+      local_state = AffineBuildState(global_state=self.affine_state,
+                                     allow_new_symbols=False)
       exprs = []
       for expr in operand_def.index_dims:
         exprs.append(expr.build(state=local_state))
       self.operands[operand_def].indexing_map = _ir.AffineMap.get(
           dim_count=local_state.dim_count,
           symbol_count=local_state.symbol_count,
           exprs=exprs)
 
   def add_tensor_use(self, tensor_use: TensorUse):
     if tensor_use in self.uses:
       return
     with self.context:
-      local_state = AffineBuildState(
-          global_state=self.affine_state, allow_new_symbols=False)
+      local_state = AffineBuildState(global_state=self.affine_state,
+                                     allow_new_symbols=False)
       exprs = []
       for expr in tensor_use.indices:
         exprs.append(expr.build(state=local_state))
-      indexing_map = _ir.AffineMap.get(
-          dim_count=local_state.dim_count,
+      indexing_map = _ir.AffineMap.get(dim_count=local_state.dim_count,
                                        symbol_count=local_state.symbol_count,
                                        exprs=exprs)
 
       use_config = TensorUseConfig(tensor_use, indexing_map)
       self.uses[tensor_use] = use_config
 
   def _get_scalar_map(self) -> _ir.AffineMap:
     """Create an empty affine map used to index a scalar."""
     with self.context:
-      return _ir.AffineMap.get(
-          dim_count=self.affine_state.dim_count,
+      return _ir.AffineMap.get(dim_count=self.affine_state.dim_count,
                                symbol_count=self.affine_state.symbol_count,
                                exprs=list())
 
   def _normalize_affine_map(self,
                             affine_map: _ir.AffineMap,
                             with_dims: bool = True) -> _ir.AffineMap:
     """Normalizes an indexing map to have the max known symbols and dims."""
     with self.context:
       return _ir.AffineMap.get(
           dim_count=self.affine_state.dim_count if with_dims else 0,
class LinalgOpConfig(YAMLObject):
   @staticmethod
   def from_linalg_op_def(
       op_def: LinalgOpDef,
       context: Optional[_ir.Context] = None) -> Sequence["LinalgOpConfig"]:
     """Expands a LinalgOpDef into corresponding Linalg configured ops."""
     # TODO: Many LinalgOpDef patterns need to expand to multiple generics.
     assert len(op_def.comprehensions) == 1, "Only one comprehension supported"
     return [
-        LinalgOpConfig(
-            op_def.metadata,
+        LinalgOpConfig(op_def.metadata,
                        structured_op=LinalgStructuredOpConfig(
                            op_def.comprehensions[0], op_def.domain,
                            op_def.registered_operands.values(), context)),
     ]
 
   def __repr__(self):
     return (f"LinalgOpConfig(metadata={self.metadata},\n"
             f"structured_op={self.structured_op})")

mlir/python/mlir/dialects/linalg/opdsl/lang/dsl.py

keyword argument is interpreted here:
         False, a named form is emitted (which must have been built in to the
         compiler).
     """
     emit_generic = kwargs.pop("emit_generic", False)
     if not isinstance(emit_generic, bool):
       raise ValueError(f"The named argument 'emit_generic' needs to be "
                        f" of type bool but got {type(emit_generic)}")
 
-    op_configs = LinalgOpConfig.from_linalg_op_def(
-        self.op_def, context=ir.Context.current)
+    op_configs = LinalgOpConfig.from_linalg_op_def(self.op_def,
+                                                   context=ir.Context.current)
 
     if len(op_configs) != 1:
       # TODO: Support composite ops.
       raise NotImplementedError(
           f"Emission of composite linalg ops not supported: {op_configs}")
 
     ctx = ir.Context.current
     linalgDialect = ctx.get_dialect_descriptor("linalg")
     fully_qualified_name = "linalg." + self.op_name
-    emit_generic = (
-        emit_generic or not ctx.is_registered_operation(fully_qualified_name))
+    emit_generic = (emit_generic or
+                    not ctx.is_registered_operation(fully_qualified_name))
 
     op_config = op_configs[0]
     out_values = _prepare_structured_op_outs(outs)
     in_values = [_get_op_result_or_value(i) for i in ins]
     if op_config.structured_op:
       if emit_generic:
-        return emit_generic_structured_op(
-            op_config.structured_op, *in_values, outs=out_values, **kwargs)
+        return emit_generic_structured_op(op_config.structured_op,
+                                          *in_values,
+                                          outs=out_values,
+                                          **kwargs)
       else:
-        return emit_named_structured_op(
-            op_config.structured_op,
+        return emit_named_structured_op(op_config.structured_op,
                                         self.op_name,
                                         self.op_def.metadata.cpp_class_name,
                                         *in_values,
                                         outs=out_values,
                                         **kwargs)
 
     raise NotImplementedError(
         f"Emission of linalg op type not supported: {op_config}")
 
 
 def linalg_structured_op(dsl_func=None,
                          *,
                          op_name=None,
                          op_class_name=None) -> DefinedOpCallable:
   if dsl_func is None:
     # Curry the keyword args in for delayed application.
-    return functools.partial(
-        linalg_structured_op, op_name=op_name, op_class_name=op_class_name)
+    return functools.partial(linalg_structured_op,
+                             op_name=op_name,
+                             op_class_name=op_class_name)
   # Determine default names by introspecting the function.
   if op_name is None:
     op_name = dsl_func.__name__
   if op_class_name is None:
     # Camel case it.
     op_class_name = f"{''.join(x.title() for x in op_name.split('_'))}Op"
 
-  op_def = LinalgOpDef(
-      name=op_name, cpp_class_name=op_class_name, doc=inspect.getdoc(dsl_func))
+  op_def = LinalgOpDef(name=op_name,
+                       cpp_class_name=op_class_name,
+                       doc=inspect.getdoc(dsl_func))
 
   # Extract arguments and TensorDefs from the signature.
   dsl_func_args = list()
   sig = inspect.signature(dsl_func)
   for param_name, param in sig.parameters.items():
     param_default = param.default
     if isinstance(param_default,
                   (TensorDef, ScalarDef, IndexAttrDef, UnaryFnAttrDef,

mlir/python/mlir/dialects/linalg/opdsl/lang/emitter.py

if expr.scalar_arg:
         return self.block_arg_mapping[expr.scalar_arg.arg]
       except KeyError:
         raise ValueError(f"Argument {expr.scalar_arg.arg} is not bound for "
                          f"this structured op.")
     elif expr.scalar_const:
       value_attr = Attribute.parse(expr.scalar_const.value)
       return arith.ConstantOp(value_attr.type, value_attr).result
     elif expr.scalar_index:
-      dim_attr = IntegerAttr.get(
-          IntegerType.get_signless(64), expr.scalar_index.dim)
+      dim_attr = IntegerAttr.get(IntegerType.get_signless(64),
+                                 expr.scalar_index.dim)
       return linalg.IndexOp(dim_attr).result
     elif expr.scalar_fn:
       kind = expr.scalar_fn.kind.name.lower()
       fn_name = expr.scalar_fn.fn_name
       if expr.scalar_fn.attr_name:
         fn_name, _ = self.fn_attr_mapping[expr.scalar_fn.attr_name]
       fn = self._get_function(f"_{kind}_{fn_name}")
       operand_values = [

mlir/python/mlir/dialects/linalg/opdsl/ops/core_named_ops.py

def pooling_nhwc_min_unsigned(I=TensorDef(T1, S.N, S.OH * S.SH + S.KH * S.DH,
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.oh, D.ow, D.c, D.kh, D.kw)
   O[D.n, D.oh, D.ow,
     D.c] = ReduceFn.min_unsigned[D.kh, D.kw](TypeFn.cast_unsigned(
         U, I[D.n, D.oh * S.SH + D.kh * S.DH, D.ow * S.SW + D.kw * S.DW, D.c]))
 
+
 @linalg_structured_op
-def pooling_nwc_sum(I=TensorDef(T1, S.N,
-                                S.OW * S.SW + S.KW * S.DW, S.C),
+def pooling_nwc_sum(I=TensorDef(T1, S.N, S.OW * S.SW + S.KW * S.DW, S.C),
                     K=TensorDef(T2, S.KW, index_dims=[D.kw]),
                     O=TensorDef(U, S.N, S.OW, S.C, output=True),
                     strides=IndexAttrDef(S.SW, default=[1]),
                     dilations=IndexAttrDef(S.DW, default=[1])):
   """Performs sum pooling.
 
   Layout:
     * Input: NWC.
     * Kernel: W.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.ow, D.c, D.kw)
-  O[D.n, D.ow, D.c] += TypeFn.cast_signed(
-      U, I[D.n, D.ow * S.SW + D.kw * S.DW, D.c])
+  O[D.n, D.ow,
+    D.c] += TypeFn.cast_signed(U, I[D.n, D.ow * S.SW + D.kw * S.DW, D.c])
 
 
 @linalg_structured_op
-def pooling_ncw_sum(I=TensorDef(T1, S.N, S.C,
-                                S.OW * S.SW + S.KW * S.DW),
+def pooling_ncw_sum(I=TensorDef(T1, S.N, S.C, S.OW * S.SW + S.KW * S.DW),
                     K=TensorDef(T2, S.KW, index_dims=[D.kw]),
                     O=TensorDef(U, S.N, S.C, S.OW, output=True),
                     strides=IndexAttrDef(S.SW, default=[1]),
                     dilations=IndexAttrDef(S.DW, default=[1])):
   """Performs sum pooling.
 
   Layout:
     * Input: NCW.
     * Kernel: W.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.c, D.ow, D.kw)
-  O[D.n, D.c, D.ow] += TypeFn.cast_signed(
-      U, I[D.n, D.c, D.ow * S.SW + D.kw * S.DW])
+  O[D.n, D.c,
+    D.ow] += TypeFn.cast_signed(U, I[D.n, D.c, D.ow * S.SW + D.kw * S.DW])
 
 
 @linalg_structured_op
-def pooling_nwc_max(I=TensorDef(T1, S.N,
-                                S.OW * S.SW + S.KW * S.DW, S.C),
+def pooling_nwc_max(I=TensorDef(T1, S.N, S.OW * S.SW + S.KW * S.DW, S.C),
                     K=TensorDef(T2, S.KW, index_dims=[D.kw]),
                     O=TensorDef(U, S.N, S.OW, S.C, output=True),
                     strides=IndexAttrDef(S.SW, default=[1]),
                     dilations=IndexAttrDef(S.DW, default=[1])):
   """Performs max pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.ow, D.c, D.kw)
   O[D.n, D.ow, D.c] = ReduceFn.max_signed[[D.kw]](TypeFn.cast_signed(
       U, I[D.n, D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
-def pooling_nwc_max_unsigned(I=TensorDef(T1, S.N,
-                                         S.OW * S.SW + S.KW * S.DW, S.C),
-                             K=TensorDef(T2,
-                                         S.KW,
-                                         index_dims=[D.kw]),
+def pooling_nwc_max_unsigned(I=TensorDef(T1, S.N, S.OW * S.SW + S.KW * S.DW,
+                                         S.C),
+                             K=TensorDef(T2, S.KW, index_dims=[D.kw]),
                              O=TensorDef(U, S.N, S.OW, S.C, output=True),
                              strides=IndexAttrDef(S.SW, default=[1]),
                              dilations=IndexAttrDef(S.DW, default=[1])):
   """Performs unsigned max pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.ow, D.c, D.kw)
-  O[D.n, D.ow,
-    D.c] = ReduceFn.max_unsigned[[D.kw]](TypeFn.cast_unsigned(
-        U, I[D.n, D.ow * S.SW + D.kw * S.DW, D.c]))
+  O[D.n, D.ow, D.c] = ReduceFn.max_unsigned[[D.kw]](TypeFn.cast_unsigned(
 
 
 @linalg_structured_op
-def pooling_ncw_max(I=TensorDef(T1, S.N, S.C,
-                                S.OW * S.SW + S.KW * S.DW),
+def pooling_ncw_max(I=TensorDef(T1, S.N, S.C, S.OW * S.SW + S.KW * S.DW),
                     K=TensorDef(T2, S.KW, index_dims=[D.kw]),
                     O=TensorDef(U, S.N, S.C, S.OW, output=True),
                     strides=IndexAttrDef(S.SW, default=[1]),
                     dilations=IndexAttrDef(S.DW, default=[1])):
   """Performs max pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.c, D.ow, D.kw)
   O[D.n, D.c, D.ow] = ReduceFn.max_signed[[D.kw]](TypeFn.cast_signed(
       U, I[D.n, D.c, D.ow * S.SW + D.kw * S.DW,]))
 
 
 @linalg_structured_op
-def pooling_nwc_min(I=TensorDef(T1, S.N,
-                                S.OW * S.SW + S.KW * S.DW, S.C),
+def pooling_nwc_min(I=TensorDef(T1, S.N, S.OW * S.SW + S.KW * S.DW, S.C),
                     K=TensorDef(T2, S.KW, index_dims=[D.kw]),
                     O=TensorDef(U, S.N, S.OW, S.C, output=True),
                     strides=IndexAttrDef(S.SW, default=[1]),
                     dilations=IndexAttrDef(S.DW, default=[1])):
   """Performs min pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.ow, D.c, D.kw)
   O[D.n, D.ow, D.c] = ReduceFn.min_signed[[D.kw]](TypeFn.cast_signed(
       U, I[D.n, D.ow * S.SW + D.kw * S.DW, D.c]))
 
 
 @linalg_structured_op
-def pooling_nwc_min_unsigned(I=TensorDef(T1, S.N,
-                                         S.OW * S.SW + S.KW * S.DW, S.C),
-                             K=TensorDef(T2,
-                                         S.KW,
-                                         index_dims=[D.kw]),
+def pooling_nwc_min_unsigned(I=TensorDef(T1, S.N, S.OW * S.SW + S.KW * S.DW,
+                                         S.C),
+                             K=TensorDef(T2, S.KW, index_dims=[D.kw]),
                              O=TensorDef(U, S.N, S.OW, S.C, output=True),
                              strides=IndexAttrDef(S.SW, default=[1]),
                              dilations=IndexAttrDef(S.DW, default=[1])):
   """Performs unsigned min pooling.
 
   Numeric casting is performed on the input operand, promoting it to the same
   data type as the accumulator/output.
   """
   implements(ConvolutionOpInterface)
   domain(D.n, D.ow, D.c, D.kw)
-  O[D.n, D.ow,
-    D.c] = ReduceFn.min_unsigned[[D.kw]](TypeFn.cast_unsigned(
-        U, I[D.n, D.ow * S.SW + D.kw * S.DW, D.c]))
+  O[D.n, D.ow, D.c] = ReduceFn.min_unsigned[[D.kw]](TypeFn.cast_unsigned(
 
 
-
 @linalg_structured_op
 def pooling_ndhwc_sum(I=TensorDef(T1, S.N, S.OD * S.SD + S.KD * S.DD,
                                   S.OH * S.SH + S.KH * S.DH,
                                   S.OW * S.SW + S.KW * S.DW, S.C),
                       K=TensorDef(T2,
                                   S.KD,
                                   S.KH,
                                   S.KW,

mlir/python/mlir/dialects/python_test.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from ._python_test_ops_gen import *
 from .._mlir_libs._mlirPythonTest import TestAttr, TestType, TestTensorValue
 
+
 def register_python_test_dialect(context, load=True):
   from .._mlir_libs import _mlirPythonTest
   _mlirPythonTest.register_python_test_dialect(context, load)

mlir/python/mlir/dialects/transform/__init__.py

class FailurePropagationMode(Enum):
 
   def _as_int(self):
     if self is FailurePropagationMode.PROPAGATE:
       return 1
 
     assert self is FailurePropagationMode.SUPPRESS
     return 2
 
+
 from .._transform_ops_gen import *
 from ..._mlir_libs._mlirDialectsTransform import *

mlir/python/mlir/execution_engine.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 # Simply a wrapper around the extension module of the same name.
 from ._mlir_libs import _mlirExecutionEngine as _execution_engine
 import ctypes
 
 __all__ = [
     "ExecutionEngine",
 ]
 
+
 class ExecutionEngine(_execution_engine.ExecutionEngine):
 
   def lookup(self, name):
     """Lookup a function emitted with the `llvm.emit_c_interface`
     attribute and returns a ctype callable.
     Raise a RuntimeError if the function isn't found.
     """
     func = self.raw_lookup("_mlir_ciface_" + name)

mlir/python/mlir/runtime/np_to_memref.py

strided_arr = np.lib.stride_tricks.as_strided(
       np.ctypeslib.as_array(val[0].shape),
       np.ctypeslib.as_array(val[0].strides) * np_arr.itemsize,
   )
   return to_numpy(strided_arr)
 
 
 def ranked_memref_to_numpy(ranked_memref):
   """Converts ranked memrefs to numpy arrays."""
-  np_arr = np.ctypeslib.as_array(
-      ranked_memref[0].aligned, shape=ranked_memref[0].shape)
+  np_arr = np.ctypeslib.as_array(ranked_memref[0].aligned,
+                                 shape=ranked_memref[0].shape)
   strided_arr = np.lib.stride_tricks.as_strided(
       np_arr,
       np.ctypeslib.as_array(ranked_memref[0].shape),
       np.ctypeslib.as_array(ranked_memref[0].strides) * np_arr.itemsize,
   )
   return to_numpy(strided_arr)