tico 0.1.0.dev250722__py3-none-any.whl → 0.1.0.dev250724__py3-none-any.whl

tico/__init__.py

@@ -20,8 +20,16 @@ from packaging.version import Version
 from tico.config import CompileConfigV1, get_default_config
 from tico.utils.convert import convert, convert_from_exported_program, convert_from_pt2
 
+__all__ = [
+    "CompileConfigV1",
+    "get_default_config",
+    "convert",
+    "convert_from_exported_program",
+    "convert_from_pt2",
+]
+
 # THIS LINE IS AUTOMATICALLY GENERATED BY setup.py
-__version__ = "0.1.0.dev250722"
+__version__ = "0.1.0.dev250724"
 
 MINIMUM_SUPPORTED_VERSION = "2.5.0"
 SECURE_TORCH_VERSION = "2.6.0"

tico/config/base.py

@@ -31,7 +31,7 @@ class CompileConfigBase:
         config = cls()
         for key in config_dict:
             if key in config.to_dict():
-                assert type(config.get(key)) == bool
+                assert isinstance(config.get(key), bool)
                 config.set(key, config_dict[key])
 
         return config

tico/experimental/quantization/__init__.py

@@ -1 +1,6 @@
 from tico.experimental.quantization.public_interface import convert, prepare
+
+__all__ = [
+    "convert",
+    "prepare",
+]

tico/experimental/quantization/algorithm/pt2e/annotation/annotator.py

@@ -21,12 +21,7 @@ if TYPE_CHECKING:
     import torch.fx
     from torch.ao.quantization.qconfig import _ObserverOrFakeQuantizeConstructor
 import torch
-from torch.ao.quantization.observer import (
-    MinMaxObserver,
-    MovingAverageMinMaxObserver,
-    MovingAveragePerChannelMinMaxObserver,
-    PerChannelMinMaxObserver,
-)
+from torch.ao.quantization.observer import MinMaxObserver, PerChannelMinMaxObserver
 from torch.ao.quantization.quantizer import QuantizationSpec, Quantizer
 from torch.ao.quantization.quantizer.utils import _get_module_name_filter
 

tico/experimental/quantization/algorithm/pt2e/annotation/op/linear.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, List, Optional, TYPE_CHECKING
+from typing import Callable, Optional, TYPE_CHECKING
 
 if TYPE_CHECKING:
     import torch.fx

tico/experimental/quantization/algorithm/pt2e/utils.py

@@ -19,7 +19,6 @@ if TYPE_CHECKING:
 import torch
 from torch.ao.quantization.quantizer import QuantizationSpec
 from torch.ao.quantization.quantizer.utils import _get_module_name_filter
-from torch.utils import _pytree as pytree
 
 from tico.experimental.quantization.algorithm.pt2e.annotation.config import (
     QuantizationConfig,

tico/experimental/quantization/algorithm/smoothquant/smooth_quant.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Dict, List, Optional
+from typing import Dict, List, Optional
 
 import torch
 

tico/experimental/quantization/evaluation/evaluate.py

@@ -166,7 +166,7 @@ def evaluate(
         )
     if not isinstance(backend, BACKEND):
         raise RuntimeError(
-            f"Invalid backend. Please use tico.quantization.evaluate.BACKEND enum class"
+            "Invalid backend. Please use tico.quantization.evaluate.BACKEND enum class"
         )
     # Make it a list for simpler logic.
     if input_data is not None:

tico/experimental/quantization/passes/fold_quant_ops.py

@@ -16,7 +16,6 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     import torch.fx
-import copy
 
 import torch
 from torch.export import ExportedProgram

tico/experimental/quantization/passes/insert_quantize_on_dtype_mismatch.py

@@ -19,7 +19,7 @@ if TYPE_CHECKING:
 import copy
 
 from collections import defaultdict
-from typing import Any, Callable
+from typing import Any
 
 import torch
 from torch.export import ExportedProgram

tico/experimental/quantization/passes/quantize_bias.py

@@ -16,7 +16,6 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     import torch.fx
-import copy
 
 import torch
 from torch.export import ExportedProgram

tico/experimental/quantization/passes/remove_weight_dequant_op.py

@@ -53,7 +53,7 @@ class ValRange:
         if isinstance(val, torch.Tensor):
             self.max = torch.max(val).item()
             self.min = torch.min(val).item()
-        elif type(val) == list:
+        elif isinstance(val, list):
             self.max = max(val)
             self.min = min(val)
         else:

tico/passes/cast_aten_where_arg_type.py

@@ -176,7 +176,7 @@ class CastATenWhereArgType(PassBase):
             node_dtype = extract_torch_dtype(node)
             assert (
                 node_dtype == node_dtype_ori
-            ), f"Type casting doesn't change node's dtype."
+            ), "Type casting doesn't change node's dtype."
 
             logger.debug(
                 f"{to_cast.name}'s dtype was casted from {buf_data.dtype} to {dtype_to_cast}"

tico/passes/cast_mixed_type_args.py

@@ -124,7 +124,7 @@ class CastMixedTypeArgs(PassBase):
             if rhs_val.dtype == type_to_promote:
                 ori_type = lhs_val.dtype
                 arg_to_promote = lhs
-            assert arg_to_promote != None
+            assert arg_to_promote is not None
 
             if isinstance(arg_to_promote, torch.fx.Node):
                 with graph.inserting_after(arg_to_promote):
@@ -178,7 +178,7 @@ class CastMixedTypeArgs(PassBase):
             node_dtype = extract_torch_dtype(node)
             assert (
                 node_dtype == node_dtype_ori
-            ), f"Type casting doesn't change node's dtype."
+            ), "Type casting doesn't change node's dtype."
 
         graph.eliminate_dead_code()
         graph.lint()

tico/passes/const_prop_pass.py

@@ -301,7 +301,7 @@ class ConstPropPass(PassBase):
         graph.eliminate_dead_code()
         graph_module.recompile()
 
-        logger.debug(f"Constant nodes are propagated")
+        logger.debug("Constant nodes are propagated")
         # Constant folding can be done with only one time run. Let's set `modified` to False.
         modified = False
         return PassResult(modified)

tico/passes/convert_conv1d_to_conv2d.py

@@ -19,7 +19,7 @@ if TYPE_CHECKING:
 import torch
 from torch.export import ExportedProgram
 
-from tico.serialize.circle_graph import extract_shape
+from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
 from tico.utils.errors import NotYetSupportedError
 from tico.utils.graph import create_node

tico/passes/decompose_addmm.py

@@ -20,7 +20,6 @@ import torch
 from torch.export import ExportedProgram
 
 from tico.serialize.circle_mapping import extract_shape
-from tico.utils import logging
 from tico.utils.graph import add_placeholder, create_node
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
@@ -59,8 +58,6 @@ class DecomposeAddmm(PassBase):
         super().__init__()
 
     def call(self, exported_program: ExportedProgram) -> PassResult:
-        logger = logging.getLogger(__name__)
-
         gm = exported_program.graph_module
         graph: torch.fx.Graph = gm.graph
         modified = False

tico/passes/decompose_batch_norm.py

@@ -96,9 +96,9 @@ class DecomposeBatchNorm(PassBase):
             eps = args.eps
 
             if not running_mean:
-                raise NotYetSupportedError(f"running_mean=None is not supported yet")
+                raise NotYetSupportedError("running_mean=None is not supported yet")
             if not running_var:
-                raise NotYetSupportedError(f"running_var=None is not supported yet")
+                raise NotYetSupportedError("running_var=None is not supported yet")
 
             """
             Only support the cases generated from torch.nn.BatchNorm2d module,

tico/passes/decompose_fake_quantize.py

@@ -19,10 +19,8 @@ if TYPE_CHECKING:
 import torch
 
 # To import torch.ops.quantized_decomposed related operator
-from torch.ao.quantization.fx._decomposed import quantized_decomposed_lib
 from torch.export import ExportedProgram
 
-from tico.utils import logging
 from tico.utils.graph import create_node
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
@@ -66,7 +64,6 @@ class DecomposeFakeQuantize(PassBase):
         super().__init__()
 
     def call(self, exported_program: ExportedProgram) -> PassResult:
-        logger = logging.getLogger(__name__)
         modified = False
 
         gm = exported_program.graph_module

tico/passes/decompose_fake_quantize_tensor_qparams.py

@@ -26,10 +26,8 @@ from torch._export.utils import (
 )
 
 # To import torch.ops.quantized_decomposed related operator
-from torch.ao.quantization.fx._decomposed import quantized_decomposed_lib
 from torch.export import ExportedProgram
 
-from tico.utils import logging
 from tico.utils.graph import create_node
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import (

tico/passes/decompose_group_norm.py

@@ -22,7 +22,6 @@ import torch
 from torch.export import ExportedProgram
 
 from tico.serialize.circle_mapping import extract_shape
-from tico.utils import logging
 from tico.utils.graph import create_node
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
@@ -126,8 +125,6 @@ class DecomposeGroupNorm(PassBase):
         )
 
     def call(self, exported_program: ExportedProgram) -> PassResult:
-        logger = logging.getLogger(__name__)
-
         gm = exported_program.graph_module
         graph: torch.fx.Graph = gm.graph
         modified = False

tico/passes/legalize_predefined_layout_operators.py

@@ -20,7 +20,7 @@ if TYPE_CHECKING:
 import torch
 from torch.export import ExportedProgram
 
-from tico.serialize.circle_graph import extract_shape
+from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
 from tico.utils.errors import NotYetSupportedError
 from tico.utils.graph import create_node
@@ -206,7 +206,6 @@ class LegalizePreDefinedLayoutOperators(PassBase):
 
         args = ConvTranspose2DArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
         input = args.input
-        padding = args.padding
         groups = args.groups
         dilation = args.dilation
 
@@ -288,13 +287,12 @@ class LegalizePreDefinedLayoutOperators(PassBase):
         input = args.input
         weight = args.weight
         bias = args.bias
-        eps = args.eps
 
         running_mean = args.running_mean
         running_var = args.running_var
         use_input_stats = args.use_input_stats
 
-        if not (use_input_stats == True):
+        if not use_input_stats:
             raise NotYetSupportedError("Only support use_input_stats is True.")
         if not isinstance(running_mean, NoneType):
             raise NotYetSupportedError("Only support running_mean=None")
@@ -350,10 +348,6 @@ class LegalizePreDefinedLayoutOperators(PassBase):
         # max_pool2d_with_indices(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> (Tensor, Tensor)
         args = MaxPool2dWithIndicesArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
         input_ = args.input
-        kernel_size = args.kernel_size
-        stride = args.stride
-        padding = args.padding
-        dilation = args.dilation
         ceil_mode = args.ceil_mode
         if ceil_mode:
             raise NotYetSupportedError("Only support non-ceil model.")
@@ -402,9 +396,6 @@ class LegalizePreDefinedLayoutOperators(PassBase):
         # avg_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, bool ceil_mode=False, bool count_include_pad=True, int? divisor_override=None) -> (Tensor)
         args = AvgPool2dArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
         input_ = args.input
-        kernel_size = args.kernel_size
-        stride = args.stride
-        padding = args.padding
         ceil_mode = args.ceil_mode
         if ceil_mode:
             raise NotYetSupportedError("Only support non-ceil model.")

tico/passes/lower_to_resize_nearest_neighbor.py

@@ -67,7 +67,7 @@ class LowerToResizeNearestNeighbor(PassBase):
             return None
         # indices = [None, None, H index, W index]
         N, C, H, W = indices
-        if N != None or C != None:
+        if N is not None or C is not None:
             return None
         if not isinstance(H, torch.fx.Node):
             return None

tico/passes/lower_to_slice.py

@@ -28,7 +28,7 @@ from torch._export.utils import (
 from torch.export import ExportedProgram
 
 from tico.passes import ops
-from tico.serialize.circle_graph import extract_shape
+from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
 from tico.utils.graph import create_node, is_single_value_tensor
 from tico.utils.passes import PassBase, PassResult

tico/passes/merge_consecutive_cat.py

@@ -51,7 +51,7 @@ class MergeConsecutiveCat(PassBase):
                 if not prev_cat.op == "call_function":
                     continue
 
-                if not prev_cat.target in ops.aten.cat:
+                if prev_cat.target not in ops.aten.cat:
                     continue
 
                 prev_args = CatArgs(*prev_cat.args, **prev_cat.kwargs)  # type: ignore[arg-type]

tico/passes/remove_redundant_expand.py

@@ -12,11 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    import torch.fx
-import torch
 from torch.export import ExportedProgram
 
 from tico.passes import ops

tico/passes/remove_redundant_reshape.py

@@ -90,7 +90,7 @@ class RemoveRedundantReshapePattern1(PassBase):
             if len(permute.users) != 1:
                 continue
             permute_args = PermuteArgs(*permute.args, **permute.kwargs)  # type: ignore[arg-type]
-            permute_input, permute_dims = permute_args.input, permute_args.dims
+            permute_dims = permute_args.dims
             # (1xAxBxC) - `aten.permute` - (1xAxCxB)
             if permute_dims != [0, 1, 3, 2]:
                 continue
@@ -172,7 +172,7 @@ class RemoveRedundantReshapePattern2(PassBase):
             if len(permute.users) != 1:
                 continue
             permute_args = PermuteArgs(*permute.args, **permute.kwargs)  # type: ignore[arg-type]
-            permute_input, permute_dims = permute_args.input, permute_args.dims
+            permute_dims = permute_args.dims
             # (1xAxBxC) - `aten.permute` - (Bx1xAxC)
             if permute_dims != [2, 0, 1, 3]:
                 continue
@@ -262,7 +262,7 @@ class RemoveRedundantReshapePattern3(PassBase):
                 continue
 
             # add
-            if not add.target in ops.aten.add:
+            if add.target not in ops.aten.add:
                 continue
             add_args = AddTensorArgs(*add.args, **add.kwargs)  # type: ignore[arg-type]
             reshape_2, reshape_3 = add_args.input, add_args.other
@@ -272,7 +272,7 @@ class RemoveRedundantReshapePattern3(PassBase):
             # reshape_2
             if not reshape_2.op == "call_function":
                 continue
-            if not reshape_2.target in ops.aten.reshape:
+            if reshape_2.target not in ops.aten.reshape:
                 continue
             reshape_2_args = ReshapeArgs(*reshape_2.args, **reshape_2.kwargs)  # type: ignore[arg-type]
             reshape_2_input = reshape_2_args.input
@@ -280,7 +280,7 @@ class RemoveRedundantReshapePattern3(PassBase):
             # reshape_3
             if not reshape_3.op == "call_function":
                 continue
-            if not reshape_3.target in ops.aten.reshape:
+            if reshape_3.target not in ops.aten.reshape:
                 continue
             reshape_3_args = ReshapeArgs(*reshape_3.args, **reshape_3.kwargs)  # type: ignore[arg-type]
             reshape_3_input = reshape_3_args.input

tico/passes/segment_index_select.py

@@ -29,7 +29,7 @@ from torch._export.utils import (
 from torch.export import ExportedProgram
 
 from tico.passes import ops
-from tico.serialize.circle_graph import extract_shape
+from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
 from tico.utils.graph import add_placeholder, create_node, is_single_value_tensor
 from tico.utils.passes import PassBase, PassResult

tico/serialize/circle_graph.py

@@ -323,7 +323,7 @@ class CircleSubgraph(circle.SubGraph.SubGraphT):
         self, node: Union[torch.fx.Node, circle.Tensor.TensorT, ConstData]
     ) -> int:
         # return -1 if node is None. This is for generating CircleOutputExclude
-        if node == None:
+        if node is None:
             return -1
 
         if hasattr(node, "name") and node.name in self.name_to_tid: