polish repeat: support higher dimension

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)

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
 

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()