Making sure indexing operation is not dropping dimensions

- Added relevant test

Author: pavanky

arrayfire/array.py

@@ -426,7 +426,7 @@ def __getitem__(self, key):
         try:
             out = Array()
             n_dims = self.numdims()
-            inds = get_indices(key, n_dims)
+            inds = get_indices(key)
 
             safe_call(clib.af_index_gen(ct.pointer(out.arr),
                                         self.arr, ct.c_longlong(n_dims), ct.pointer(inds)))
@@ -446,7 +446,7 @@ def __setitem__(self, key, val):
                 other_arr = val.arr
 
             out_arr = ct.c_void_p(0)
-            inds  = get_indices(key, n_dims)
+            inds  = get_indices(key)
 
             safe_call(clib.af_assign_gen(ct.pointer(out_arr),
                                          self.arr, ct.c_longlong(n_dims), ct.pointer(inds),

arrayfire/index.py

@@ -104,13 +104,11 @@ def __init__ (self, idx):
         else:
             self.idx.seq = Seq(idx)
 
-def get_indices(key, n_dims):
+def get_indices(key):
 
-    index_vec = Index * n_dims
-    inds = index_vec()
-
-    for n in range(n_dims):
-        inds[n] = Index(slice(None))
+    index_vec = Index * 4
+    S = Index(slice(None))
+    inds = index_vec(S, S, S, S)
 
     if isinstance(key, tuple):
         n_idx = len(key)
@@ -143,9 +141,6 @@ def get_assign_dims(key, idims):
     elif isinstance(key, tuple):
         n_inds = len(key)
 
-        if (n_inds > len(idims)):
-            raise IndexError("Number of indices greater than array dimensions")
-
         for n in range(n_inds):
             if (is_number(key[n])):
                 dims[n] = 1

tests/simple_index.py

@@ -58,3 +58,11 @@
 for ii in ParallelRange(2,5):
     b[ii] = 2
 af.display(b)
+
+a = af.randu(3,2)
+rows = af.constant(0, 1, dtype=af.s32)
+b = a[:,rows]
+af.display(b)
+for r in rows:
+    af.display(r)
+    af.display(b[:,r])