From 2070aaae09b409c9801b1d326bdc22be5241a1a6 Mon Sep 17 00:00:00 2001
From: Megvii Engine Team <megengine@megvii.com>
Date: Thu, 24 Sep 2020 16:20:40 +0800
Subject: [PATCH] fix(mge/functional): fix duplicate `transpose` in `__all__`
 and add `dimshuffle

GitOrigin-RevId: 29670882338fd43d4c91519e1acaaa6e454bb992
---
 imperative/python/megengine/functional/tensor.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/imperative/python/megengine/functional/tensor.py b/imperative/python/megengine/functional/tensor.py
index 4cdebc155..8586f0b69 100644
--- a/imperative/python/megengine/functional/tensor.py
+++ b/imperative/python/megengine/functional/tensor.py
@@ -37,8 +37,7 @@ __all__ = [
     "broadcast",
     "concat",
     "cond_take",
-    "transpose",
-    "add_axis",
+    "dimshuffle",
     "eye",
     "flatten",
     "full",
@@ -732,7 +731,10 @@ def transpose(inp: Tensor, pattern: Iterable[int]) -> Tensor:
     return inp.transpose(pattern)
 
 
-dimshuffle = transpose
+def dimshuffle(inp: Tensor, pattern: Iterable[int]) -> Tensor:
+    r"""Same as :func:`~.transpose`.
+    """
+    return transpose(inp, pattern)
 
 
 def reshape(inp: Tensor, target_shape: Iterable[int]) -> Tensor:
-- 
GitLab