[arrayfire] 361/408: Restore original shape after flattening input for linear indexing

[Ghislain Vaillant]

This is an automated email from the git hooks/post-receive script.

ghisvail-guest pushed a commit to branch debian/sid
in repository arrayfire.

commit 654dffbab92cedb588ea07490d4ef38f81a1b819
Author: Pavan Yalamanchili <pavan at arrayfire.com>
Date:   Wed Aug 26 12:00:25 2015 -0400

    Restore original shape after flattening input for linear indexing
    
    - Added checks for the same
---
 src/api/c/assign.cpp | 23 ++++++++++++++---------
 test/assign.cpp      |  9 +++++++++
 2 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/src/api/c/assign.cpp b/src/api/c/assign.cpp
index 7d5e651..3c93ca2 100644
--- a/src/api/c/assign.cpp
+++ b/src/api/c/assign.cpp
@@ -121,12 +121,15 @@ af_err af_assign_seq(af_array *out,
         ARG_ASSERT(1, (ndims>0));
         ARG_ASSERT(3, (rhs!=0));
 
+        ArrayInfo lInfo = getInfo(lhs);
 
-        if (ndims == 1 && ndims != (dim_t)getInfo(lhs).ndims()) {
-            af_array tmp;
-            AF_CHECK(af_flat(&tmp, lhs));
-            AF_CHECK(af_assign_seq(out, tmp, ndims, index, rhs));
-            AF_CHECK(af_release_array(tmp));
+        if (ndims == 1 && ndims != (dim_t)lInfo.ndims()) {
+            af_array tmp_in, tmp_out;
+            AF_CHECK(af_flat(&tmp_in, lhs));
+            AF_CHECK(af_assign_seq(&tmp_out, tmp_in, ndims, index, rhs));
+            AF_CHECK(af_moddims(out, tmp_out, lInfo.ndims(), lInfo.dims().get()));
+            AF_CHECK(af_release_array(tmp_in));
+            AF_CHECK(af_release_array(tmp_out));
             return AF_SUCCESS;
         }
 
@@ -228,10 +231,12 @@ af_err af_assign_gen(af_array *out,
         ARG_ASSERT(2, (ndims == 1) || (ndims == (dim_t)lInfo.ndims()));
 
         if (ndims == 1 && ndims != (dim_t)lInfo.ndims()) {
-            af_array tmp;
-            AF_CHECK(af_flat(&tmp, lhs));
-            AF_CHECK(af_assign_gen(out, tmp, ndims, indexs, rhs_));
-            AF_CHECK(af_release_array(tmp));
+            af_array tmp_in, tmp_out;
+            AF_CHECK(af_flat(&tmp_in, lhs));
+            AF_CHECK(af_assign_gen(&tmp_out, tmp_in, ndims, indexs, rhs_));
+            AF_CHECK(af_moddims(out, tmp_out, lInfo.ndims(), lInfo.dims().get()));
+            AF_CHECK(af_release_array(tmp_in));
+            AF_CHECK(af_release_array(tmp_out));
             return AF_SUCCESS;
         }
 
diff --git a/test/assign.cpp b/test/assign.cpp
index 3abe1b4..5692392 100644
--- a/test/assign.cpp
+++ b/test/assign.cpp
@@ -846,6 +846,9 @@ TEST(Asssign, LinearCPP)
     af::index idx = af::seq(st, en);
     a(idx) = 3;
 
+    ASSERT_EQ(a.dims(0), a_copy.dims(0));
+    ASSERT_EQ(a.dims(1), a_copy.dims(1));
+
     std::vector<float> ha(nx * ny);
     std::vector<float> ha_copy(nx * ny);
 
@@ -884,6 +887,8 @@ TEST(Asssign, LinearAssignSeq)
 
     af::array out(out_arr);
 
+    ASSERT_EQ(a.dims(0), out.dims(0));
+    ASSERT_EQ(a.dims(1), out.dims(1));
 
     std::vector<float> hout(nx * ny);
     std::vector<float> ha(nx * ny);
@@ -923,6 +928,8 @@ TEST(Asssign, LinearAssignGenSeq)
 
     af::array out(out_arr);
 
+    ASSERT_EQ(a.dims(0), out.dims(0));
+    ASSERT_EQ(a.dims(1), out.dims(1));
 
     std::vector<float> hout(nx * ny);
     std::vector<float> ha(nx * ny);
@@ -962,6 +969,8 @@ TEST(Asssign, LinearAssignGenArr)
 
     af::array out(out_arr);
 
+    ASSERT_EQ(a.dims(0), out.dims(0));
+    ASSERT_EQ(a.dims(1), out.dims(1));
 
     std::vector<float> hout(nx * ny);
     std::vector<float> ha(nx * ny);

-- 
Alioth's /usr/local/bin/git-commit-notice on /srv/git.debian.org/git/debian-science/packages/arrayfire.git