From e5217db9f9dadee5ad686d8663834fb650ab3946 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Fri, 27 Aug 2021 15:22:40 +0800
Subject: [PATCH] feat(traced_module): move traced_module out of the
 experimental folder

GitOrigin-RevId: 36c76b5277c64b15ef4a9aacdf7fa4e1eb936c10
---
 imperative/python/megengine/__init__.py       |  1 -
 .../python/megengine/experimental/__init__.py |  1 -
 .../traced_module/__init__.py                 |  2 +-
 .../{experimental => }/traced_module/expr.py  | 15 ++++++-----
 .../traced_module/fake_quant.py               | 14 ++++++++---
 .../traced_module/module_tracer.py            | 11 ++++----
 .../{experimental => }/traced_module/node.py  |  7 +++---
 .../traced_module/pytree.py                   | 15 ++++++-----
 .../traced_module/serialization.py            |  6 ++---
 .../traced_module/traced_module.py            | 25 +++++++++----------
 .../{experimental => }/traced_module/utils.py |  3 +--
 .../python/test/integration/test_converge.py  |  2 +-
 .../test_converge_with_gradient_clip.py       |  2 +-
 .../test/integration/test_trace_dump.py       |  2 +-
 .../test/unit/core/test_serialization.py      |  5 +---
 .../python/test/unit/module/test_module.py    |  2 +-
 .../test/unit/traced_module/test_haoruitao.py |  2 +-
 .../test/unit/traced_module/test_jit_trace.py |  2 +-
 .../unit/traced_module/test_modification.py   |  4 +--
 .../unit/traced_module/test_serialization.py  |  2 +-
 .../unit/traced_module/test_trace_module.py   |  2 +-
 .../test/unit/traced_module/test_wujianan.py  |  2 +-
 22 files changed, 62 insertions(+), 65 deletions(-)
 rename imperative/python/megengine/{experimental => }/traced_module/__init__.py (89%)
 rename imperative/python/megengine/{experimental => }/traced_module/expr.py (97%)
 rename imperative/python/megengine/{experimental => }/traced_module/fake_quant.py (73%)
 rename imperative/python/megengine/{experimental => }/traced_module/module_tracer.py (97%)
 rename imperative/python/megengine/{experimental => }/traced_module/node.py (98%)
 rename imperative/python/megengine/{experimental => }/traced_module/pytree.py (95%)
 rename imperative/python/megengine/{experimental => }/traced_module/serialization.py (89%)
 rename imperative/python/megengine/{experimental => }/traced_module/traced_module.py (99%)
 rename imperative/python/megengine/{experimental => }/traced_module/utils.py (99%)

diff --git a/imperative/python/megengine/__init__.py b/imperative/python/megengine/__init__.py
index f149bcfde..ef453fc74 100644
--- a/imperative/python/megengine/__init__.py
+++ b/imperative/python/megengine/__init__.py
@@ -130,4 +130,3 @@ import megengine.optimizer
 import megengine.quantization
 import megengine.random
 import megengine.utils
-import megengine.experimental
diff --git a/imperative/python/megengine/experimental/__init__.py b/imperative/python/megengine/experimental/__init__.py
index d263b902c..19b1fc6bf 100644
--- a/imperative/python/megengine/experimental/__init__.py
+++ b/imperative/python/megengine/experimental/__init__.py
@@ -6,5 +6,4 @@
 # Unless required by applicable law or agreed to in writing,
 # software distributed under the License is distributed on an
 # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-from . import traced_module
 from .weight_scaler import get_scaled_model
diff --git a/imperative/python/megengine/experimental/traced_module/__init__.py b/imperative/python/megengine/traced_module/__init__.py
similarity index 89%
rename from imperative/python/megengine/experimental/traced_module/__init__.py
rename to imperative/python/megengine/traced_module/__init__.py
index 1a34651ba..970225bf0 100644
--- a/imperative/python/megengine/experimental/traced_module/__init__.py
+++ b/imperative/python/megengine/traced_module/__init__.py
@@ -6,7 +6,7 @@
 # software distributed under the License is distributed on an
 # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 
-from ...core._imperative_rt.core2 import set_cpp_apply_module_trace
+from ..core._imperative_rt.core2 import set_cpp_apply_module_trace
 from .traced_module import (
     TracedModule,
     _register_all_builtin_module,
diff --git a/imperative/python/megengine/experimental/traced_module/expr.py b/imperative/python/megengine/traced_module/expr.py
similarity index 97%
rename from imperative/python/megengine/experimental/traced_module/expr.py
rename to imperative/python/megengine/traced_module/expr.py
index a5d119c5d..e42ca2417 100644
--- a/imperative/python/megengine/experimental/traced_module/expr.py
+++ b/imperative/python/megengine/traced_module/expr.py
@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 # MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 #
 # Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
@@ -14,13 +13,13 @@ import inspect
 import re
 from typing import Callable, Dict, List
 
-from ...core._imperative_rt import OpDef
-from ...core._imperative_rt.core2 import Tensor as RawTensor
-from ...core._imperative_rt.core2 import apply, set_module_tracing, unset_module_tracing
-from ...core.ops.builtin import FakeQuant
-from ...core.ops.special import Const
-from ...module import Module
-from ...tensor import Parameter, Tensor
+from ..core._imperative_rt import OpDef
+from ..core._imperative_rt.core2 import Tensor as RawTensor
+from ..core._imperative_rt.core2 import apply, set_module_tracing, unset_module_tracing
+from ..core.ops.builtin import FakeQuant
+from ..core.ops.special import Const
+from ..module import Module
+from ..tensor import Parameter, Tensor
 from .module_tracer import active_module_tracer, module_tracer
 from .node import ModuleNode, Node, NodeMixin, TensorNode
 from .pytree import ArgsIndex, TreeDef, _is_const_leaf, _is_leaf, tree_flatten
diff --git a/imperative/python/megengine/experimental/traced_module/fake_quant.py b/imperative/python/megengine/traced_module/fake_quant.py
similarity index 73%
rename from imperative/python/megengine/experimental/traced_module/fake_quant.py
rename to imperative/python/megengine/traced_module/fake_quant.py
index d81d3c94f..120cdc6b1 100644
--- a/imperative/python/megengine/experimental/traced_module/fake_quant.py
+++ b/imperative/python/megengine/traced_module/fake_quant.py
@@ -1,9 +1,17 @@
+# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
+#
+# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+
 from copy import deepcopy
 from typing import Union
 
-from ...core.tensor.dtype import QuantDtypeMeta
-from ...quantization.fake_quant import QParamsModuleMixin, _FakeQuantize
-from ...quantization.utils import QParams, QuantMode, fake_quant_tensor
+from ..core.tensor.dtype import QuantDtypeMeta
+from ..quantization.fake_quant import QParamsModuleMixin, _FakeQuantize
+from ..quantization.utils import QParams, QuantMode, fake_quant_tensor
 
 
 class FakeQuantize(_FakeQuantize, QParamsModuleMixin):
diff --git a/imperative/python/megengine/experimental/traced_module/module_tracer.py b/imperative/python/megengine/traced_module/module_tracer.py
similarity index 97%
rename from imperative/python/megengine/experimental/traced_module/module_tracer.py
rename to imperative/python/megengine/traced_module/module_tracer.py
index 49632d291..6505fe724 100644
--- a/imperative/python/megengine/experimental/traced_module/module_tracer.py
+++ b/imperative/python/megengine/traced_module/module_tracer.py
@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 # MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 #
 # Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
@@ -8,11 +7,11 @@
 # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 import collections
 
-from ... import Tensor
-from ... import functional as F
-from ...core.tensor.array_method import ArrayMethodMixin
-from ...module import Module
-from ...module.qat import QATModule
+from .. import Tensor
+from .. import functional as F
+from ..core.tensor.array_method import ArrayMethodMixin
+from ..module import Module
+from ..module.qat import QATModule
 
 _active_module_tracer = None
 
diff --git a/imperative/python/megengine/experimental/traced_module/node.py b/imperative/python/megengine/traced_module/node.py
similarity index 98%
rename from imperative/python/megengine/experimental/traced_module/node.py
rename to imperative/python/megengine/traced_module/node.py
index 15c89e01d..6e0330584 100644
--- a/imperative/python/megengine/experimental/traced_module/node.py
+++ b/imperative/python/megengine/traced_module/node.py
@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 # MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 #
 # Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
@@ -12,9 +11,9 @@ from typing import Any, Dict, List, Tuple, Type
 
 import numpy
 
-from ...core._imperative_rt.core2 import Tensor as RawTensor
-from ...module import Module
-from ...tensor import Tensor
+from ..core._imperative_rt.core2 import Tensor as RawTensor
+from ..module import Module
+from ..tensor import Tensor
 
 
 class Node:
diff --git a/imperative/python/megengine/experimental/traced_module/pytree.py b/imperative/python/megengine/traced_module/pytree.py
similarity index 95%
rename from imperative/python/megengine/experimental/traced_module/pytree.py
rename to imperative/python/megengine/traced_module/pytree.py
index 686b651e8..e4f13e4bc 100644
--- a/imperative/python/megengine/experimental/traced_module/pytree.py
+++ b/imperative/python/megengine/traced_module/pytree.py
@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 # MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 #
 # Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
@@ -13,13 +12,13 @@ from typing import Callable, NamedTuple
 
 import numpy as np
 
-from ...core._imperative_rt.common import CompNode
-from ...core._imperative_rt.core2 import Tensor as RawTensor
-from ...core._wrap import Device
-from ...core.tensor.dtype import QuantDtypeMeta
-from ...module import Module
-from ...quantization.utils import LSQParams, QParams, QuantMode
-from ...tensor import Parameter, Tensor
+from ..core._imperative_rt.common import CompNode
+from ..core._imperative_rt.core2 import Tensor as RawTensor
+from ..core._wrap import Device
+from ..core.tensor.dtype import QuantDtypeMeta
+from ..module import Module
+from ..quantization.utils import LSQParams, QParams, QuantMode
+from ..tensor import Parameter, Tensor
 from .node import ModuleNode, Node, NodeMixin, TensorNode
 
 
diff --git a/imperative/python/megengine/experimental/traced_module/serialization.py b/imperative/python/megengine/traced_module/serialization.py
similarity index 89%
rename from imperative/python/megengine/experimental/traced_module/serialization.py
rename to imperative/python/megengine/traced_module/serialization.py
index ec9596216..8ce3ed67a 100644
--- a/imperative/python/megengine/experimental/traced_module/serialization.py
+++ b/imperative/python/megengine/traced_module/serialization.py
@@ -7,9 +7,9 @@
 # "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 from typing import Dict
 
-from ...core._imperative_rt import OpDef
-from ...core.ops import builtin
-from ...version import __version__
+from ..core._imperative_rt import OpDef
+from ..core.ops import builtin
+from ..version import __version__
 
 OPDEF_PARAM_LOADER = {}
 
diff --git a/imperative/python/megengine/experimental/traced_module/traced_module.py b/imperative/python/megengine/traced_module/traced_module.py
similarity index 99%
rename from imperative/python/megengine/experimental/traced_module/traced_module.py
rename to imperative/python/megengine/traced_module/traced_module.py
index 15f93552c..4bd9553ce 100644
--- a/imperative/python/megengine/experimental/traced_module/traced_module.py
+++ b/imperative/python/megengine/traced_module/traced_module.py
@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 # MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 #
 # Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
@@ -22,21 +21,21 @@ from typing import Callable, Dict, Iterable, List, Optional, Sequence, Type, Uni
 
 from megengine import tensor
 
-from ... import functional as F
-from ... import get_logger
-from ... import module as M
-from ...core._imperative_rt.core2 import Tensor as RawTensor
-from ...core._imperative_rt.core2 import (
+from .. import functional as F
+from .. import get_logger
+from .. import module as M
+from ..core._imperative_rt.core2 import Tensor as RawTensor
+from ..core._imperative_rt.core2 import (
     is_tracing_module,
     set_module_tracing,
     unset_module_tracing,
 )
-from ...core._trace_option import set_symbolic_shape
-from ...core.tensor.array_method import ArrayMethodMixin
-from ...module import Module
-from ...module.qat import QATModule
-from ...quantization.fake_quant import LSQ, TQT, FakeQuantize, _FakeQuantize
-from ...quantization.observer import (
+from ..core._trace_option import set_symbolic_shape
+from ..core.tensor.array_method import ArrayMethodMixin
+from ..module import Module
+from ..module.qat import QATModule
+from ..quantization.fake_quant import LSQ, TQT, FakeQuantize, _FakeQuantize
+from ..quantization.observer import (
     ExponentialMovingAverageObserver,
     HistogramObserver,
     MinMaxObserver,
@@ -45,7 +44,7 @@ from ...quantization.observer import (
     SyncExponentialMovingAverageObserver,
     SyncMinMaxObserver,
 )
-from ...tensor import Tensor
+from ..tensor import Tensor
 from .expr import Apply, CallFunction, CallMethod, Constant, Expr, GetAttr, Input
 from .fake_quant import FakeQuantize as TM_FakeQuant
 from .module_tracer import (
diff --git a/imperative/python/megengine/experimental/traced_module/utils.py b/imperative/python/megengine/traced_module/utils.py
similarity index 99%
rename from imperative/python/megengine/experimental/traced_module/utils.py
rename to imperative/python/megengine/traced_module/utils.py
index dffb7da28..c7c52e73b 100644
--- a/imperative/python/megengine/experimental/traced_module/utils.py
+++ b/imperative/python/megengine/traced_module/utils.py
@@ -1,4 +1,3 @@
-# -*- coding: utf-8 -*-
 # MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
 #
 # Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
@@ -10,7 +9,7 @@ import copy
 from collections.abc import MutableMapping, MutableSequence
 from typing import Dict, Iterable, List, Optional, Sequence
 
-from ...module import Module
+from ..module import Module
 
 
 def replace_container_with_module_container(container):
diff --git a/imperative/python/test/integration/test_converge.py b/imperative/python/test/integration/test_converge.py
index fb080b414..fc93542fa 100644
--- a/imperative/python/test/integration/test_converge.py
+++ b/imperative/python/test/integration/test_converge.py
@@ -15,9 +15,9 @@ import megengine as mge
 import megengine.autodiff as ad
 import megengine.functional as F
 from megengine import Tensor
-from megengine.experimental.traced_module import trace_module
 from megengine.module import Linear, Module
 from megengine.optimizer import SGD
+from megengine.traced_module import trace_module
 
 batch_size = 64
 data_shape = (batch_size, 2)
diff --git a/imperative/python/test/integration/test_converge_with_gradient_clip.py b/imperative/python/test/integration/test_converge_with_gradient_clip.py
index 9a0b83935..5a91c5981 100644
--- a/imperative/python/test/integration/test_converge_with_gradient_clip.py
+++ b/imperative/python/test/integration/test_converge_with_gradient_clip.py
@@ -16,10 +16,10 @@ import megengine.autodiff as ad
 import megengine.functional as F
 import megengine.optimizer as optim
 from megengine import Tensor
-from megengine.experimental.traced_module import trace_module
 from megengine.jit import trace
 from megengine.module import Linear, Module
 from megengine.optimizer import SGD
+from megengine.traced_module import trace_module
 
 batch_size = 64
 data_shape = (batch_size, 2)
diff --git a/imperative/python/test/integration/test_trace_dump.py b/imperative/python/test/integration/test_trace_dump.py
index e256354a4..085cc396a 100644
--- a/imperative/python/test/integration/test_trace_dump.py
+++ b/imperative/python/test/integration/test_trace_dump.py
@@ -19,8 +19,8 @@ import megengine.module as M
 import megengine.optimizer as optim
 from megengine import tensor
 from megengine.autodiff import GradManager
-from megengine.experimental.traced_module import trace_module
 from megengine.jit import trace
+from megengine.traced_module import trace_module
 
 
 @contextlib.contextmanager
diff --git a/imperative/python/test/unit/core/test_serialization.py b/imperative/python/test/unit/core/test_serialization.py
index 7eae4ea3e..382efe8d4 100644
--- a/imperative/python/test/unit/core/test_serialization.py
+++ b/imperative/python/test/unit/core/test_serialization.py
@@ -15,10 +15,7 @@ import numpy as np
 import megengine as mge
 from megengine import Parameter, Tensor
 from megengine.core.ops import builtin
-from megengine.experimental.traced_module.serialization import (
-    get_opdef_state,
-    load_opdef_from_state,
-)
+from megengine.traced_module.serialization import get_opdef_state, load_opdef_from_state
 
 
 def test_tensor_serialization():
diff --git a/imperative/python/test/unit/module/test_module.py b/imperative/python/test/unit/module/test_module.py
index 083b2a527..b7f789046 100644
--- a/imperative/python/test/unit/module/test_module.py
+++ b/imperative/python/test/unit/module/test_module.py
@@ -15,7 +15,6 @@ import pytest
 import megengine as mge
 import megengine.functional as F
 from megengine import Parameter, Tensor, tensor
-from megengine.experimental.traced_module import TracedModule, trace_module
 from megengine.module import (
     BatchNorm1d,
     BatchNorm2d,
@@ -30,6 +29,7 @@ from megengine.module import (
 )
 from megengine.module.module import _access_structure
 from megengine.quantization.quantize import quantize, quantize_qat
+from megengine.traced_module import TracedModule, trace_module
 from megengine.utils.module_utils import get_expand_structure, set_expand_structure
 
 
diff --git a/imperative/python/test/unit/traced_module/test_haoruitao.py b/imperative/python/test/unit/traced_module/test_haoruitao.py
index bdf8cd4a1..c371a409b 100644
--- a/imperative/python/test/unit/traced_module/test_haoruitao.py
+++ b/imperative/python/test/unit/traced_module/test_haoruitao.py
@@ -7,8 +7,8 @@ import megengine.functional as F
 import megengine.module as M
 import megengine.utils.comp_graph_tools as cgtools
 from megengine.core._trace_option import set_symbolic_shape
-from megengine.experimental.traced_module import trace_module
 from megengine.jit import trace
+from megengine.traced_module import trace_module
 
 set_symbolic_shape(True)
 
diff --git a/imperative/python/test/unit/traced_module/test_jit_trace.py b/imperative/python/test/unit/traced_module/test_jit_trace.py
index 0cd3bcae5..fa76cf96f 100644
--- a/imperative/python/test/unit/traced_module/test_jit_trace.py
+++ b/imperative/python/test/unit/traced_module/test_jit_trace.py
@@ -12,9 +12,9 @@ import numpy as np
 import megengine.functional as F
 import megengine.module as M
 import megengine.utils.comp_graph_tools as cgtools
-from megengine.experimental.traced_module import trace_module
 from megengine.jit import trace
 from megengine.module import Module
+from megengine.traced_module import trace_module
 
 
 class MyBlock(Module):
diff --git a/imperative/python/test/unit/traced_module/test_modification.py b/imperative/python/test/unit/traced_module/test_modification.py
index 593a54487..5b3312858 100644
--- a/imperative/python/test/unit/traced_module/test_modification.py
+++ b/imperative/python/test/unit/traced_module/test_modification.py
@@ -9,8 +9,8 @@ import numpy as np
 
 import megengine.functional as F
 import megengine.module as M
-from megengine.experimental.traced_module import trace_module
-from megengine.experimental.traced_module.expr import CallFunction, GetAttr
+from megengine.traced_module import trace_module
+from megengine.traced_module.expr import CallFunction, GetAttr
 
 
 class MyBlock(M.Module):
diff --git a/imperative/python/test/unit/traced_module/test_serialization.py b/imperative/python/test/unit/traced_module/test_serialization.py
index be4d2ff00..167dd46de 100644
--- a/imperative/python/test/unit/traced_module/test_serialization.py
+++ b/imperative/python/test/unit/traced_module/test_serialization.py
@@ -12,8 +12,8 @@ import numpy as np
 import megengine.functional as F
 import megengine.module as M
 from megengine import Tensor
-from megengine.experimental.traced_module import trace_module
 from megengine.module import Module
+from megengine.traced_module import trace_module
 
 
 class MyBlock(Module):
diff --git a/imperative/python/test/unit/traced_module/test_trace_module.py b/imperative/python/test/unit/traced_module/test_trace_module.py
index 5ccf935de..fd0d8f610 100644
--- a/imperative/python/test/unit/traced_module/test_trace_module.py
+++ b/imperative/python/test/unit/traced_module/test_trace_module.py
@@ -2,7 +2,7 @@ import numpy as np
 
 import megengine.module as M
 from megengine import Tensor
-from megengine.experimental.traced_module import TracedModule, trace_module
+from megengine.traced_module import TracedModule, trace_module
 
 
 class MyModule1(M.Module):
diff --git a/imperative/python/test/unit/traced_module/test_wujianan.py b/imperative/python/test/unit/traced_module/test_wujianan.py
index 44474f6da..fb87761fa 100644
--- a/imperative/python/test/unit/traced_module/test_wujianan.py
+++ b/imperative/python/test/unit/traced_module/test_wujianan.py
@@ -8,8 +8,8 @@ import megengine.functional as F
 import megengine.module as M
 import megengine.utils.comp_graph_tools as cgtools
 from megengine.core._trace_option import set_symbolic_shape
-from megengine.experimental.traced_module import trace_module
 from megengine.jit import trace
+from megengine.traced_module import trace_module
 
 set_symbolic_shape(True)
 
-- 
GitLab