Linear indexing now flattens the arrays before the operation

pavanky · pavanky · commit 72952ed74ddd · 2015-03-12T13:59:56.000-04:00
diff --git a/src/api/cpp/array.cpp b/src/api/cpp/array.cpp
@@ -299,6 +299,14 @@ namespace af
     array array::operator()(const array& idx) const
     {
         eval();
+
+        // Special case of indexing linearly
+        // Flatten the current array and index accordingly
+        if (this->numdims() > 1) {
+            array tmp = flat(*this);
+            return tmp(idx);
+        }
+
         af_array out = 0;
         AF_THROW(af_lookup(&out, this->get(), idx.get(), 0));
         return array(out);
@@ -307,6 +315,14 @@ namespace af
     array array::operator()(const seq &s0) const
     {
         eval();
+
+        // Special case of indexing linearly
+        // Flatten the current array and index accordingly
+        if (this->numdims() > 1) {
+            array tmp = flat(*this);
+            return tmp(s0);
+        }
+
         af_array out = 0;
         seq indices[] = {s0, span, span, span};
         //FIXME: check if this->s has same dimensions as numdims
diff --git a/test/index.cpp b/test/index.cpp
@@ -29,6 +29,7 @@ using std::endl;
 using std::ostream_iterator;
 using af::dtype_traits;
 
+
 template<typename T, typename OP>
 void
 checkValues(const af_seq &seq, const T* data, const T* indexed_data, OP compair_op) {
@@ -627,7 +628,7 @@ TEST(lookup, CPP)
 
     array input(dims0, &(in[0].front()));
     array indices(dims1, &(in[1].front()));
-    array output = input(indices);
+    array output = af::lookup(input, indices, 0);
 
     vector<float> currGoldBar = tests[0];
     size_t nElems = currGoldBar.size();
@@ -728,7 +729,7 @@ TEST(SeqIndex, CPPLarge)
 
     array input(dims0, &(in[0].front()));
     array indices(dims1, &(in[1].front()));
-    array output = input(indices);
+    array output = af::lookup(input, indices, 0);
 
     vector<float> currGoldBar = tests[0];
     size_t nElems = currGoldBar.size();
