Fix lint

Author: gramalingam

examples/custom_op_expansion.py

@@ -2,15 +2,19 @@
 # Licensed under the MIT License.
 """A utility and an example showing how onnxscript functions can be used to define function expansions
 and be used with the inliner to replace calls to the custom function with an expanded subgraph.
-This is useful to perform certain classes of graph surgery easily."""
+This is useful to perform certain classes of graph surgery easily.
+"""
 
 import onnx
-import onnxscript
-from onnxscript import script, FLOAT, opset22 as op
 
+import onnxscript
 
+script = onnxscript.script
+FLOAT = onnxscript.FLOAT
+op = onnxscript.values.opset22
 local = onnxscript.values.Opset("local", 1)
 
+
 # Example Model: Actual models can come from ModelBuilder or Exporter or any other source.
 # Models can contain calls to custom operations (from a custom domain like 'local' here or
 # even "com.microsoft" etc.)
@@ -24,17 +28,20 @@ def model_script(X: FLOAT["N"], Y: FLOAT["N"]) -> FLOAT["N"]:
     Temp2 = local.CustomOp2(Temp1, alp=0.9)
     return Temp2
 
+
 # Define expansions for custom operations as onnxscript functions
 @script(opset=local)
 def CustomOp1(X: FLOAT["N"], Y: FLOAT["N"]) -> FLOAT["N"]:
     Temp1 = op.Sub(X, Y)
     return op.Div(Temp1, X)
 
+
 @script(opset=local)
 def CustomOp2(X: FLOAT["N"], alp: float) -> FLOAT["N"]:
     Temp2 = op.Elu(X, alpha=alp)
     return op.Mul(Temp2, Temp2)
 
+
 # Now, we can replace the custom operations in the model with their expansions:
 
 functions = [CustomOp1.to_function_proto(), CustomOp2.to_function_proto()]
@@ -45,8 +52,8 @@ def CustomOp2(X: FLOAT["N"], alp: float) -> FLOAT["N"]:
 print(onnx.printer.to_text(model))
 
 import onnxscript.utils.replace as replace
+
 updated_model = replace.replace_functions(model, functions)
 
 print("\nUpdated Model after replacing custom operations with their expansions:")
 print(onnx.printer.to_text(updated_model))
-

onnxscript/utils/replace.py

@@ -3,18 +3,23 @@
 """A utility function to replace custom operations in a model with their expansions"""
 
 from typing import Sequence
+
 import onnx
 import onnx_ir as ir
 import onnx_ir.passes.common as common_passes
 
-def replace_functions(model: onnx.ModelProto, functions: Sequence[onnx.FunctionProto]) -> onnx.ModelProto:
-    '''A utility function to replace custom operations in a model with their expansions:
+
+def replace_functions(
+    model: onnx.ModelProto, functions: Sequence[onnx.FunctionProto]
+) -> onnx.ModelProto:
+    """A utility function to replace custom operations in a model with their expansions:
     Args:
         model: An ONNX ModelProto possibly containing calls to custom operations.
         functions: A sequence of FunctionProto defining the expansions for the custom operations.
+
     Returns:
         An updated ModelProto with custom operations replaced by their expansions.
-    '''
+    """
     irmodel = ir.from_proto(model)
     irfunctions = [ir.from_proto(func) for func in functions]
     model_functions = irmodel.functions