From 40ce5e85cd5cfb198e6a3e9a3288e831655c1b60 Mon Sep 17 00:00:00 2001
From: lzhengning <lzhengning@gmail.com>
Date: Tue, 7 Sep 2021 20:15:33 +0800
Subject: [PATCH] polish repeat: support higher dimension

---
 python/jittor/misc.py             | 10 ++---
 python/jittor/test/test_loss3d.py | 10 +++++
 python/jittor/test/test_repeat.py | 61 +++++++++++++++++++++++++++++++
 3 files changed, 76 insertions(+), 5 deletions(-)
 create mode 100644 python/jittor/test/test_repeat.py

diff --git a/python/jittor/misc.py b/python/jittor/misc.py
index 3045b15f..6cd5cf53 100644
--- a/python/jittor/misc.py
+++ b/python/jittor/misc.py
@@ -100,15 +100,15 @@ def repeat(x, *shape):
         x_shape = (len_shape - len_x_shape) * [1] + x.shape
         x = x.broadcast(x_shape)
     elif len_x_shape > len_shape:
-        rep_shape = (len_x_shape - len_shape) * [1] + shape
-    #TODO if input.shape[i]=1, no add [1]
+        rep_shape = (len_x_shape - len_shape) * [1] + list(shape)
+
     reshape_shape = []
     broadcast_shape = []
     for x_s,r_s in zip(x_shape,rep_shape):
-        reshape_shape.append(1)
+        if r_s != 1:
+            reshape_shape.append(1)
+            broadcast_shape.append(r_s)
         reshape_shape.append(x_s)
-
-        broadcast_shape.append(r_s)
         broadcast_shape.append(1)
 
     x = x.reshape(reshape_shape)
diff --git a/python/jittor/test/test_loss3d.py b/python/jittor/test/test_loss3d.py
index 7178871d..864a4e78 100644
--- a/python/jittor/test/test_loss3d.py
+++ b/python/jittor/test/test_loss3d.py
@@ -1,3 +1,13 @@
+# ***************************************************************
+# Copyright (c) 2021 Jittor. All Rights Reserved. 
+# Maintainers: 
+#     Zheng-Ning Liu <lzhengning@gmail.com>
+#     Dun Liang <randonlang@gmail.com>. 
+# 
+# This file is subject to the terms and conditions defined in
+# file 'LICENSE.txt', which is part of this source code package.
+# ***************************************************************
+
 import unittest
 import numpy as np
 
diff --git a/python/jittor/test/test_repeat.py b/python/jittor/test/test_repeat.py
new file mode 100644
index 00000000..15ca7fd9
--- /dev/null
+++ b/python/jittor/test/test_repeat.py
@@ -0,0 +1,61 @@
+# ***************************************************************
+# Copyright (c) 2021 Jittor. All Rights Reserved. 
+# Maintainers: 
+#     Zheng-Ning Liu <lzhengning@gmail.com>
+#     Dun Liang <randonlang@gmail.com>. 
+# 
+# This file is subject to the terms and conditions defined in
+# file 'LICENSE.txt', which is part of this source code package.
+# ***************************************************************
+
+
+import unittest
+import jittor as jt
+import numpy as np
+
+
+class TestRepeatOp(unittest.TestCase):
+    def test_repeat(self):
+        np_a = np.arange(5)
+        jt_a = jt.array(np_a)
+
+        np_b = np.tile(np_a, (2, 3))
+        jt_b = jt.repeat(jt_a, (2, 3))
+
+        assert np.allclose(np_b, jt_b.data)
+
+        np_b = np.tile(np_a, (2, 3, 1))
+        jt_b = jt.repeat(jt_a, (2, 3, 1))
+
+        assert np.allclose(np_b, jt_b.data)
+
+        np_a = np.arange(24).reshape(2, 3, 4)
+        jt_a = jt.array(np_a)
+
+        np_b = np.tile(np_a, (2, 3))
+        jt_b = jt.repeat(jt_a, (2, 3))
+
+        assert np.allclose(np_b, jt_b.data)
+
+
+    def test_highdim(self):
+        np_a = np.arange(64).reshape(2, 2, 2, 2, 2, 2)
+        jt_a = jt.array(np_a)
+
+        np_b = np.tile(np_a, (2, 3))
+        jt_b = jt.repeat(jt_a, (2, 3))
+
+        assert np.allclose(np_b, jt_b.data)
+
+        np_b = np.tile(np_a, (2, 1, 1, 3))
+        jt_b = jt.repeat(jt_a, (2, 1, 1, 3))
+
+        assert np.allclose(np_b, jt_b.data)
+
+        np_b = np.tile(np_a, (2, 1, 1, 1, 3, 1))
+        jt_b = jt.repeat(jt_a, (2, 1, 1, 1, 3, 1))
+
+        assert np.allclose(np_b, jt_b.data)
+
+if __name__ == "__main__":
+    unittest.main()
\ No newline at end of file