From fdd9516e982dad030b2fac36acab4184e27e532b Mon Sep 17 00:00:00 2001
From: Marcin Junczys-Dowmunt <junczys@amu.edu.pl>
Date: Fri, 27 Oct 2017 00:07:20 +0200
Subject: fixed bug in transpose for ndarray

---
 src/graph/node_operators_unary.h | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

(limited to 'src/graph/node_operators_unary.h')

diff --git a/src/graph/node_operators_unary.h b/src/graph/node_operators_unary.h
index 852d6d02..9881357c 100644
--- a/src/graph/node_operators_unary.h
+++ b/src/graph/node_operators_unary.h
@@ -702,9 +702,12 @@ struct TransposeNodeOp : public UnaryNodeOp {
 
   template <class... Args>
   Shape newShape(Expr a, Shape permute) {
-    Shape shape;
+    Shape shape = a->shape();
+
+    UTIL_THROW_IF2(shape.size() != permute.size(),
+                   "Shape and transpose axis have different number of dimensions");
 
-    for(int i = 0; i < 4; ++i)
+    for(int i = 0; i < shape.size(); ++i)
       shape.set(i, a->shape()[permute[i]]);
 
     return shape;
-- 
cgit v1.2.3