Merge pull request #54 from tharapalanivel/fx_lint

Lint for fx

Author: andrea-fasoli

fms_mo/fx/__init__.py

@@ -0,0 +1,16 @@
+# Copyright The FMS Model Optimizer Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Imports from fx
+"""

fms_mo/fx/dynamo_utils.py

@@ -17,6 +17,7 @@
 """
 
 # Standard
+from typing import Dict, Optional
 import logging
 
 # Third Party
@@ -46,7 +47,7 @@ def dfs_gm(
     prescreenOp=None,
     hook=None,
     return_nodes=False,
-    lut_fx_mod_name_to_org={},
+    lut_fx_mod_name_to_org: Optional[Dict[str, str]] = None,
 ):
     """Depth-First Search at FX IR level, to replace our old TorchScript equivalent func
     Because FX IR is a higher level IR, should have much fewer
@@ -227,7 +228,9 @@ def _dfs(curr_node, depth):
     return node_found
 
 
-def find_conv_on_shortcut_gm(gm: torch.fx.GraphModule, lut_fx_mod_name_to_org={}):
+def find_conv_on_shortcut_gm(
+    gm: torch.fx.GraphModule, lut_fx_mod_name_to_org: Optional[Dict[str, str]] = None
+):
     """Identify Conv on shortcut using FX GM DFS
     It's (almost) specific for ResNet-like CNNs, will return a list of module names (as used in the
     original model, not FX module names)
@@ -352,7 +355,7 @@ def find_1st_last_gm(
     firstOps=None,
     lastOps=None,
     return_1st_last_sep=False,
-    lut_fx_mod_name_to_org={},
+    lut_fx_mod_name_to_org: Optional[Dict[str, str]] = None,
 ):
     """Identify the first and last layer of interests
     Usually only interested in Conv and Linear, but could be others as well
@@ -403,7 +406,11 @@ def find_1st_last_gm(
 
 
 def find_single_sided_op_gm(
-    gm, op_of_interest=None, return_LUTs=False, verbose=False, lut_fx_mod_name_to_org={}
+    gm,
+    op_of_interest=None,
+    return_LUTs=False,
+    verbose=False,
+    lut_fx_mod_name_to_org: Optional[Dict[str, str]] = None,
 ):
     """Try to determine the "polarity" of output of "every nodes" based on their inputs and the Op
     itself, then decide which Conv/Linear (or user-specified Op) will use single-sided quantizer
@@ -576,7 +583,9 @@ def find_single_sided_op_gm(
     return SingleSidedOps
 
 
-def find_qkvsync_candidates_gm(gm, return_nodes=False, lut_fx_mod_name_to_org={}):
+def find_qkvsync_candidates_gm(
+    gm, return_nodes=False, lut_fx_mod_name_to_org: Optional[Dict[str, str]] = None
+):
     """Identify groups of Linears that share the same parent. It's a transformer-specific feature.
 
     NOTE:
@@ -626,7 +635,7 @@ def find_qkvsync_candidates_gm(gm, return_nodes=False, lut_fx_mod_name_to_org={}
     return my_1st_sibling
 
 
-def find_silu_gm(gm, lut_fx_mod_name_to_org={}):
+def find_silu_gm(gm, lut_fx_mod_name_to_org: Optional[Dict[str, str]] = None):
     """Special handle for Conv following silu, specific for EffDet and etc
     LLM could use SiLU as well (llama?), but not relavent to this func
     """
@@ -646,7 +655,12 @@ def find_silu_gm(gm, lut_fx_mod_name_to_org={}):
     return siluConv
 
 
-def find_rpn_fpn_gm(gm, verbose=False, Nsubgraph=0, lut_fx_mod_name_to_org={}):
+def find_rpn_fpn_gm(
+    gm,
+    verbose=False,
+    Nsubgraph=0,
+    lut_fx_mod_name_to_org: Optional[Dict[str, str]] = None,
+):
     """For object detection CNN models, RPN (RegionProposalNetwork) and FPN (FeaturePyramidNetwork)
     are commonly used. prefer to skip them, but may be ok to quantize in some cases.
 
@@ -750,7 +764,7 @@ def find_rpn_fpn_gm(gm, verbose=False, Nsubgraph=0, lut_fx_mod_name_to_org={}):
     return fpn_convs
 
 
-def find_and_prep_bmm_gm(gm, lut_fx_mod_name_to_org={}):
+def find_and_prep_bmm_gm(gm, lut_fx_mod_name_to_org: Optional[Dict[str, str]] = None):
     """Previously with TorchScript, we use this func to perform 2 tasks:
         a) create QBmms, and then attach them to the model,
         b) set up qcfg["which2patch_contextmanager"] so that patch_torch_bmm() context
@@ -1153,7 +1167,9 @@ def cus_backend_model_analyzer(
                 )
                 setattr(mod_bmm_happened, f"QBmm{ln}", newQBmm)
 
-    # c) identify RPN/FPN TODO this hack only works for torchvision models. will use find_rpn_fpn_gm()
+    # c) identify RPN/FPN
+    # TODO this hack only works for torchvision models. will use find_rpn_fpn_gm()
+
     # Third Party
     from torchvision.models.detection.rpn import RegionProposalNetwork
     from torchvision.ops import FeaturePyramidNetwork

fms_mo/fx/utils.py

@@ -15,7 +15,7 @@
 """Utils for FX graph parsing and external kernel lowering"""
 
 # Standard
-from typing import Any
+from typing import Any, Dict, Optional
 import logging
 import operator
 import os
@@ -317,7 +317,9 @@ def lname_to_org_name(Lname):
     return org_mod_name
 
 
-def get_org_mod_name_of_fx_node(node, gm=None, lut_fx2org={}):
+def get_org_mod_name_of_fx_node(
+    node, gm=None, lut_fx2org: Optional[Dict[str, str]] = None
+):
     """Given a FX node, could be call_module or call_fuction, find out the original module name,
     based on meta data
 
@@ -489,7 +491,7 @@ def plot_graph_module(
     skip_nodes=None,
     Nnode_to_plot=None,
     additional_coloring_rules=None,
-    lut_fx_mod_name_to_org={},
+    lut_fx_mod_name_to_org: Optional[Dict[str, str]] = None,
 ):
     """Plots a GraphModule in .SVG format to visualize the compute graph. If graphviz/pygraphviz is
     not installed properly, this function will just print out a message and do nothing.

fms_mo/utils/__init__.py

@@ -12,5 +12,5 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """
-imports all from prep.py
+Imports for utils
 """