This is an automated email from the ASF dual-hosted git repository.
junrushao pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 19c378dbe6 [Unity] Avoid Using Wildcard Imports (#15850)
19c378dbe6 is described below
commit 19c378dbe6c2e5c6fd39ea2ff6f3a31fb751a271
Author: Siyuan Feng <[email protected]>
AuthorDate: Mon Oct 2 11:39:10 2023 +0800
[Unity] Avoid Using Wildcard Imports (#15850)
Wildcard importing is not recommended in Python projects:
https://docs.python.org/3/tutorial/modules.html#importing-from-a-package
I acknowledge that disabling it will bring extra effort to adding new
operators, analysis and transforms. But it will make the code clear, also
benefit the API reference doc.
---
python/tvm/relax/analysis/__init__.py | 27 +++-
python/tvm/relax/distributed/transform/__init__.py | 3 +-
python/tvm/relax/op/__init__.py | 142 +++++++++++++++++----
python/tvm/relax/op/builtin/__init__.py | 3 +-
python/tvm/relax/op/ccl/__init__.py | 3 +-
python/tvm/relax/op/distributed/__init__.py | 3 +-
python/tvm/relax/op/grad/__init__.py | 1 -
python/tvm/relax/op/image/__init__.py | 3 +-
python/tvm/relax/op/memory/__init__.py | 3 +-
python/tvm/relax/op/nn/__init__.py | 27 +++-
python/tvm/relax/op/vm/__init__.py | 3 +-
python/tvm/relax/transform/__init__.py | 58 ++++++++-
python/tvm/script/ir_builder/relax/__init__.py | 4 +-
13 files changed, 228 insertions(+), 52 deletions(-)
diff --git a/python/tvm/relax/analysis/__init__.py
b/python/tvm/relax/analysis/__init__.py
index 7ba56ff408..cc0a36622e 100644
--- a/python/tvm/relax/analysis/__init__.py
+++ b/python/tvm/relax/analysis/__init__.py
@@ -14,8 +14,31 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=wildcard-import, redefined-builtin
"""Relax IR analysis. """
-from .analysis import *
+from .analysis import (
+ BaseCheckResult,
+ all_global_vars,
+ all_vars,
+ bound_vars,
+ contains_impure_call,
+ defined_symbolic_vars,
+ derive_call_ret_struct_info,
+ detect_recursion,
+ erase_to_well_defined,
+ free_symbolic_vars,
+ free_vars,
+ get_static_type,
+ get_var2val,
+ has_reshape_pattern,
+ name_to_binding,
+ post_order_visit,
+ remove_all_unused,
+ struct_info_base_check,
+ struct_info_lca,
+ suggest_layout_transforms,
+ tir_vars_in_struct_info,
+ udchain,
+ well_formed,
+)
from .estimate_memory_usage import estimate_memory_usage
diff --git a/python/tvm/relax/distributed/transform/__init__.py
b/python/tvm/relax/distributed/transform/__init__.py
index 125433a342..9d1e780030 100644
--- a/python/tvm/relax/distributed/transform/__init__.py
+++ b/python/tvm/relax/distributed/transform/__init__.py
@@ -14,7 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=wildcard-import, redefined-builtin
"""Relax distributed-related transformations. """
-from .transform import *
+from .transform import PropagateSharding
diff --git a/python/tvm/relax/op/__init__.py b/python/tvm/relax/op/__init__.py
index 1fa875c22e..44952efe68 100644
--- a/python/tvm/relax/op/__init__.py
+++ b/python/tvm/relax/op/__init__.py
@@ -14,40 +14,130 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=wildcard-import, redefined-builtin
+# pylint: disable= redefined-builtin
"""Relax core operators."""
-# Operators
-from .base import *
-from .binary import *
-from .create import *
-from .datatype import *
-from .index import *
-from .linear_algebra import *
-from .manipulate import *
-from .mask import *
-from .op_attrs import *
-from .statistical import *
-from .search import *
-from .set import *
-from .ternary import *
-from .unary import *
-from . import builtin
-from . import distributed
-from . import grad
-from . import image
-from . import memory
-from . import nn
-from . import ccl
-
# Register operator gradient functions
-from . import _op_gradient
+from . import _op_gradient, builtin, ccl, distributed, grad, image, memory,
nn, op_attrs
+
+# Operators
+from .base import (
+ assert_op,
+ call_builtin_with_ctx,
+ call_dps_packed,
+ call_pure_packed,
+ call_tir,
+ call_tir_inplace,
+ call_tir_with_grad,
+ hint_on_device,
+ invoke_closure,
+ invoke_pure_closure,
+ make_closure,
+ null_value,
+ print,
+ register_gradient,
+ shape_of,
+ shape_to_tensor,
+ tensor_to_shape,
+ to_vdevice,
+)
+from .binary import (
+ add,
+ bitwise_and,
+ bitwise_or,
+ bitwise_xor,
+ divide,
+ equal,
+ floor_divide,
+ greater,
+ greater_equal,
+ less,
+ less_equal,
+ logical_and,
+ logical_or,
+ logical_xor,
+ maximum,
+ minimum,
+ multiply,
+ not_equal,
+ power,
+ subtract,
+)
+from .create import (
+ arange,
+ full,
+ full_like,
+ ones,
+ ones_like,
+ tril,
+ triu,
+ zeros,
+ zeros_like,
+)
+from .datatype import astype, wrap_param
+from .index import dynamic_strided_slice, strided_slice, take
+from .linear_algebra import einsum, linear, matmul
+from .manipulate import (
+ broadcast_to,
+ collapse_sum_like,
+ collapse_sum_to,
+ concat,
+ expand_dims,
+ flatten,
+ flip,
+ layout_transform,
+ permute_dims,
+ repeat,
+ reshape,
+ scatter_elements,
+ split,
+ squeeze,
+ tile,
+)
+from .mask import masked_fill
+from .search import argmax, argmin, where
+from .set import unique
+from .statistical import cumsum, max, mean, min, prod, std, sum, variance
+from .ternary import ewise_fma
+from .unary import (
+ abs,
+ acos,
+ acosh,
+ asin,
+ asinh,
+ atan,
+ atanh,
+ bitwise_not,
+ ceil,
+ clip,
+ cos,
+ cosh,
+ erf,
+ exp,
+ floor,
+ isfinite,
+ isinf,
+ isnan,
+ log,
+ logical_not,
+ negative,
+ round,
+ rsqrt,
+ sigmoid,
+ sign,
+ sin,
+ sinh,
+ sqrt,
+ square,
+ tan,
+ tanh,
+)
def _register_op_make():
# pylint: disable=import-outside-toplevel
- from . import _ffi_api
from .. import expr
+ from . import _ffi_api
expr._op_ffi_api = _ffi_api # type: ignore
diff --git a/python/tvm/relax/op/builtin/__init__.py
b/python/tvm/relax/op/builtin/__init__.py
index 04837724b1..250abb2167 100644
--- a/python/tvm/relax/op/builtin/__init__.py
+++ b/python/tvm/relax/op/builtin/__init__.py
@@ -14,7 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=wildcard-import, redefined-builtin
"""Relax builtin operators."""
-from .builtin import *
+from .builtin import alloc_tensor, stop_lift_params
diff --git a/python/tvm/relax/op/ccl/__init__.py
b/python/tvm/relax/op/ccl/__init__.py
index 20746eb053..9ef2e5a244 100644
--- a/python/tvm/relax/op/ccl/__init__.py
+++ b/python/tvm/relax/op/ccl/__init__.py
@@ -14,6 +14,5 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=wildcard-import
"""CCL related operators."""
-from .ccl import *
+from .ccl import allgather, allreduce, broadcast_from_worker0,
scatter_from_worker0
diff --git a/python/tvm/relax/op/distributed/__init__.py
b/python/tvm/relax/op/distributed/__init__.py
index 72e6da1935..35226b8790 100644
--- a/python/tvm/relax/op/distributed/__init__.py
+++ b/python/tvm/relax/op/distributed/__init__.py
@@ -14,7 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=wildcard-import, redefined-builtin
"""Operators serving for distributed Relax."""
-from .distributed import *
+from .distributed import annotate_sharding, redistribute
diff --git a/python/tvm/relax/op/grad/__init__.py
b/python/tvm/relax/op/grad/__init__.py
index 844b8ac381..a3aa3c86db 100644
--- a/python/tvm/relax/op/grad/__init__.py
+++ b/python/tvm/relax/op/grad/__init__.py
@@ -14,7 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=wildcard-import, redefined-builtin
"""Operators serving for finding gradient of relax operators."""
from .grad import *
diff --git a/python/tvm/relax/op/image/__init__.py
b/python/tvm/relax/op/image/__init__.py
index f2552ad6ac..10ef635cbf 100644
--- a/python/tvm/relax/op/image/__init__.py
+++ b/python/tvm/relax/op/image/__init__.py
@@ -14,6 +14,5 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=wildcard-import
"""Image operators."""
-from .image import *
+from .image import resize2d
diff --git a/python/tvm/relax/op/memory/__init__.py
b/python/tvm/relax/op/memory/__init__.py
index e039590251..45819f4cb3 100644
--- a/python/tvm/relax/op/memory/__init__.py
+++ b/python/tvm/relax/op/memory/__init__.py
@@ -14,7 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=wildcard-import, redefined-builtin
"""Relax memory primitives."""
-from .memory import *
+from .memory import alloc_storage, alloc_tensor, kill_storage, kill_tensor
diff --git a/python/tvm/relax/op/nn/__init__.py
b/python/tvm/relax/op/nn/__init__.py
index af2aa106bc..d1569e11cb 100644
--- a/python/tvm/relax/op/nn/__init__.py
+++ b/python/tvm/relax/op/nn/__init__.py
@@ -14,6 +14,29 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=wildcard-import
"""Neural network related operators."""
-from .nn import *
+from .nn import (
+ adaptive_avg_pool2d,
+ attention,
+ avg_pool2d,
+ batch_norm,
+ conv1d,
+ conv1d_transpose,
+ conv2d,
+ conv2d_transpose,
+ cross_entropy_with_logits,
+ dropout,
+ gelu,
+ gelu_tanh,
+ group_norm,
+ layer_norm,
+ leakyrelu,
+ log_softmax,
+ max_pool2d,
+ nll_loss,
+ pad,
+ relu,
+ rms_norm,
+ silu,
+ softmax,
+)
diff --git a/python/tvm/relax/op/vm/__init__.py
b/python/tvm/relax/op/vm/__init__.py
index ecb2857a89..e68ecbcebd 100644
--- a/python/tvm/relax/op/vm/__init__.py
+++ b/python/tvm/relax/op/vm/__init__.py
@@ -14,7 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=wildcard-import, redefined-builtin
"""Relax vm primitives."""
-from .vm import *
+from .vm import alloc_storage, alloc_tensor, call_tir_dyn, kill_object
diff --git a/python/tvm/relax/transform/__init__.py
b/python/tvm/relax/transform/__init__.py
index 7bbe6d52c2..5dfdec8ea0 100644
--- a/python/tvm/relax/transform/__init__.py
+++ b/python/tvm/relax/transform/__init__.py
@@ -14,14 +14,64 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=wildcard-import, redefined-builtin
"""Relax transformations. """
-from .transform import *
+from .transform import (
+ AllocateWorkspace,
+ AlterOpImpl,
+ AnnotateTIROpPattern,
+ AttachGlobalSymbol,
+ BindParams,
+ BindSymbolicVars,
+ BundleModelParams,
+ CallTIRRewrite,
+ CanonicalizeBindings,
+ CombineParallelMatmul,
+ ConvertLayout,
+ DataflowBlockPass,
+ DeadCodeElimination,
+ DecomposeOpsForInference,
+ DecomposeOpsForTraining,
+ EliminateCommonSubexpr,
+ FewShotTuning,
+ FoldConstant,
+ FoldDataflowBlockOutput,
+ FunctionPass,
+ FuseOps,
+ FuseOpsByPattern,
+ FuseTIR,
+ FusionPattern,
+ Gradient,
+ KillAfterLastUse,
+ LambdaLift,
+ LegalizeOps,
+ LiftTransformParams,
+ LowerAllocTensor,
+ MergeCompositeFunctions,
+ MetaScheduleApplyDatabase,
+ MetaScheduleTuneIRMod,
+ MetaScheduleTuneTIR,
+ Normalize,
+ PatternCheckContext,
+ RealizeVDevice,
+ RemovePurityChecking,
+ RewriteCUDAGraph,
+ RewriteDataflowReshape,
+ RunCodegen,
+ SplitCallTIRByPattern,
+ StaticPlanBlockMemory,
+ ToMixedPrecision,
+ ToNonDataflow,
+ UpdateVDevice,
+ VMBuiltinLower,
+ VMShapeLower,
+ dataflowblock_pass,
+ function_pass,
+)
+
from .lazy_transform_params import LazyTransformParams
from .optimize_layout_transform import OptimizeLayoutTransform
from .remove_redundant_reshape import RemoveRedundantReshape
# Import to register the legalization functions.
-from . import legalize_ops
-from . import tuning_api
+from . import legalize_ops, tuning_api
diff --git a/python/tvm/script/ir_builder/relax/__init__.py
b/python/tvm/script/ir_builder/relax/__init__.py
index 8b495ac08d..55cb34e145 100644
--- a/python/tvm/script/ir_builder/relax/__init__.py
+++ b/python/tvm/script/ir_builder/relax/__init__.py
@@ -14,8 +14,6 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=unused-import
"""Package tvm.script.ir_builder.relax"""
-from . import frame
+from . import distributed, frame
from .ir import * # pylint: disable=wildcard-import,redefined-builtin
-from . import distributed # pylint: disable=wildcard-import,redefined-builtin
