add api.types and mark py.typed

Author: LoSealL

mppq/__init__.py

@@ -7,4 +7,7 @@
 Minimized PPQ quantizer package.
 """
 
+from mppq import api
+
+__all__ = ["api"]
 __version__ = "0.7.3"

mppq/api/__init__.py

@@ -10,8 +10,14 @@
 推荐用户仅使用该包内的函数和对象。
 """
 
-from mppq import __version__
-from mppq.api.extension import register_operation, register_platform
+from mppq.api.extension import (
+    BaseQuantizer,
+    GraphBuilder,
+    GraphDispatcher,
+    GraphExporter,
+    register_operation,
+    register_platform,
+)
 from mppq.api.interface import (
     dispatch_graph,
     export_config,
@@ -23,15 +29,11 @@
     load_quantizer,
     quantize,
 )
-from mppq.dispatcher.base import GraphDispatcher
-from mppq.executor.torch import TorchExecutor
+from mppq.api.types import TorchExecutor
 from mppq.ffi import ENABLE_CUDA_KERNEL
-from mppq.frontend.base import GraphBuilder, GraphExporter
 from mppq.logger import set_level as set_log_level
-from mppq.quantizer.base import BaseQuantizer
 
 __all__ = [
-    "__version__",
     "TorchExecutor",
     "ENABLE_CUDA_KERNEL",
     "set_log_level",

mppq/api/extension.py

@@ -11,11 +11,22 @@
 from typing import Dict, Optional, Type
 
 from mppq.dispatcher.base import DISPATCHER_TABLE, GraphDispatcher
-from mppq.executor.base import OPERATION_FORWARD_TABLE
-from mppq.executor.op.base import DEFAULT_BACKEND_TABLE, OperationForwardProtocol
+from mppq.executor.base import OPERATION_FORWARD_TABLE, BaseGraphExecutor
+from mppq.executor.op.base import (
+    ASSERT_IS_QUANT_OP,
+    ASSERT_NUM_OF_INPUT,
+    DEFAULT_BACKEND_TABLE,
+    FORCE_CONVERT_DEVICE,
+    GET_ATTRIBUTE_FROM_OPERATION,
+    GET_VALUE_FROM_INPUTS,
+    VALUE_TO_EXECUTING_DEVICE,
+    OperationForwardProtocol,
+)
 from mppq.frontend import EXPORTER, PARSER
-from mppq.ir.base.graph import GraphBuilder, GraphExporter
+from mppq.frontend.onnx.onnx_exporter import OP_CONVERTERS
+from mppq.ir.base.graph import GraphBuilder, GraphExporter, OperationExporter
 from mppq.logger import error, warning
+from mppq.quantization.optim.base import OPTIM_ALGORITHMS, QuantizationOptimizationPass
 from mppq.quantizer.base import QUANTIZER, BaseQuantizer
 
 _PLATFORM_TO_PARSER_: Dict[int, str] = {}
@@ -160,3 +171,31 @@ def register_platform(
                 f"by default the last one is used. Current quantizer is {name}"
             )
         _PLATFORM_TO_QUANTIZER_[platform_id] = name
+
+
+__all__ = [
+    "register_operation",
+    "register_platform",
+    # registry
+    "DISPATCHER_TABLE",
+    "PARSER",
+    "EXPORTER",
+    "QUANTIZER",
+    "OP_CONVERTERS",
+    "OPTIM_ALGORITHMS",
+    # base classes
+    "GraphDispatcher",
+    "GraphBuilder",
+    "GraphExporter",
+    "OperationExporter",
+    "BaseQuantizer",
+    "BaseGraphExecutor",
+    "QuantizationOptimizationPass",
+    # extension api
+    "ASSERT_IS_QUANT_OP",
+    "ASSERT_NUM_OF_INPUT",
+    "FORCE_CONVERT_DEVICE",
+    "GET_ATTRIBUTE_FROM_OPERATION",
+    "GET_VALUE_FROM_INPUTS",
+    "VALUE_TO_EXECUTING_DEVICE",
+]

mppq/api/interface.py

@@ -18,6 +18,7 @@
     _PLATFORM_TO_QUANTIZER_,
 )
 from mppq.common import DEFAULT_QUANTIZE_OP
+from mppq.defs import empty_ppq_cache
 from mppq.dispatcher import DISPATCHER_TABLE
 from mppq.dispatcher.base import GraphDispatcher
 from mppq.dispatcher.scope import IgnoredScope
@@ -27,8 +28,11 @@
 from mppq.ir.base.graph import BaseGraph, GraphBuilder, GraphExporter
 from mppq.ir.morph import GraphFormatter, GraphMerger, GraphReplacer
 from mppq.quant import TargetPrecision
+from mppq.quantization.analyse import graphwise_error_analyse, layerwise_error_analyse
 from mppq.quantizer import QUANTIZER, BaseQuantizer
 from mppq.quantizer.base import QuantizationOptimizationPass
+from mppq.utils.qfunction import ppq_fake_quant, ppq_quant_toint
+from mppq.utils.round import ppq_tensor_round
 
 
 def load_graph(graph_file: Any, parser: GraphBuilder) -> BaseGraph:
@@ -339,3 +343,22 @@ def quantize(
     )
     exporter = EXPORTER[_PLATFORM_TO_PARSER_[platform]]()
     export_graph(graph, f, exporter=exporter)
+
+
+__all__ = [
+    "load_graph",
+    "load_onnx_graph",
+    "export_graph",
+    "export_onnx_graph",
+    "export_config",
+    "format_graph",
+    "dispatch_graph",
+    "load_quantizer",
+    "quantize",
+    "empty_ppq_cache",
+    "graphwise_error_analyse",
+    "layerwise_error_analyse",
+    "ppq_tensor_round",
+    "ppq_fake_quant",
+    "ppq_quant_toint",
+]

mppq/api/types.py

@@ -0,0 +1,61 @@
+"""
+Copyright Wenyi Tang 2025
+
+:Author: Wenyi Tang
+:Email: wenyitang@outlook.com
+
+mPPQ enumerations and data types.
+"""
+
+from mppq.data import DataType
+from mppq.dispatcher import (
+    AggressiveDispatcher,
+    AllinDispatcher,
+    ConservativeDispatcher,
+    Perseus,
+    PointDispatcher,
+)
+from mppq.dispatcher.scope import IgnoredScope
+from mppq.executor import (
+    BaseGraphExecutor,
+    GraphInput,
+    QuantRuntimeHook,
+    RuntimeHook,
+    TorchExecutor,
+    TorchQuantizeDelegator,
+)
+from mppq.executor.op.base import OperationForwardProtocol, TorchBackendContext
+from mppq.ir.base.command import (
+    GraphCommand,
+    GraphCommandType,
+    GraphDeployCommand,
+    QuantizeOperationCommand,
+    ReplaceOperationCommand,
+    ReplaceVariableCommand,
+    TruncateGraphCommand,
+)
+from mppq.ir.base.graph import BaseGraph
+from mppq.ir.base.opdef import Operation, Opset, OpSocket, Variable
+from mppq.ir.base.quantize import (
+    BaseQuantFunction,
+    QuantableOperation,
+    QuantableVariable,
+)
+from mppq.ir.deploy import QuantableGraph, RunnableGraph
+from mppq.ir.morph import GraphFormatter, GraphMerger, GraphReplacer
+from mppq.ir.search import OperationSet, SearchableGraph
+from mppq.ir.training import TrainableGraph
+from mppq.quant import (
+    OperationQuantizationConfig,
+    QuantizationPolicy,
+    QuantizationProperty,
+    QuantizationStates,
+    RoundingPolicy,
+    TargetPrecision,
+    TensorQuantizationConfig,
+)
+from mppq.quantization.algorithm.training import BlockBuilder, TrainableBlock
+from mppq.quantization.optim.base import (
+    QuantizationOptimizationPass,
+    QuantizationOptimizationPipeline,
+)

mppq/executor/__init__.py

@@ -1,8 +1,14 @@
-from mppq.executor.base import BaseGraphExecutor, QuantRuntimeHook, RuntimeHook
+from mppq.executor.base import (
+    BaseGraphExecutor,
+    GraphInput,
+    QuantRuntimeHook,
+    RuntimeHook,
+)
 from mppq.executor.torch import TorchExecutor, TorchQuantizeDelegator
 
 __all__ = [
     "BaseGraphExecutor",
+    "GraphInput",
     "QuantRuntimeHook",
     "RuntimeHook",
     "TorchExecutor",

mppq/py.typed

@@ -47,7 +47,7 @@ exclude = ["mppq/executor/op/default.py"]
 
 [tool.flake8]
 ignore = ['E203', 'E231', 'E241', 'W503']
-per-file-ignores = ['__init__.py:F401']
+per-file-ignores = ['mppq/api/types.py:F401']
 max-line-length = 999
 max-complexity = 30
 count = true