adding support for array shape () passed as an index argument in advanced indexing

cunumeric/deferred.py

@@ -446,6 +446,10 @@ def _zip_indices(
         # NumPy array.
         N = self.ndim
         pointN_dtype = self.runtime.get_point_type(N)
+        # if scalar array is passed as an argument, make the output
+        # shape be (1,)
+        if out_shape == ():
+            out_shape = (1,)
         store = self.context.create_store(
             pointN_dtype, shape=out_shape, optimize_scalar=True
         )
@@ -463,7 +467,8 @@ def _zip_indices(
         task.add_scalar_arg(self.shape, (ty.int64,))
         for a in arrays:
             task.add_input(a)
-            task.add_alignment(output_arr.base, a)
+            if a.shape != ():
+                task.add_alignment(output_arr.base, a)
         task.execute()
 
         return output_arr

src/cunumeric/index/zip_template.inl

@@ -67,12 +67,12 @@ struct ZipImpl {
 template <VariantKind KIND>
 static void zip_template(TaskContext& context)
 {
-  // Here `N` is the number of dimenstions of the input array and the number
+  // Here `N` is the number of dimensions of the input array and the number
   // of dimensions of the Point<N> field
   // key_dim - is the number of dimensions of the index arrays before
   // they were broadcasted to the shape of the input array (shape of
   // all index arrays should be the same))
-  // start index - is the index from wich first index array was passed
+  // start index - is the index from which first index array was passed
   // DIM - dimension of the output array
   //
   // for the example:
@@ -95,7 +95,11 @@ static void zip_template(TaskContext& context)
   int64_t start_index = context.scalars()[2].value<int64_t>();
   auto shape          = context.scalars()[3].value<DomainPoint>();
   ZipArgs args{context.outputs()[0], context.inputs(), N, key_dim, start_index, shape};
-  double_dispatch(args.inputs[0].dim(), N, ZipImpl<KIND>{}, args);
+  int dim = args.inputs[0].dim();
+  // if scalar passed as an input, convert it to the array size 1
+  if (dim == 0) { dim = 1; }
+
+  double_dispatch(dim, N, ZipImpl<KIND>{}, args);
 }
 
 }  // namespace cunumeric

tests/integration/test_set_item.py

@@ -40,9 +40,9 @@ def test_basic():
 def test_scalar_ndarray_as_index(arr):
     offsets = num.arange(5)  # [0, 1, 2, 3, 4]
     offset = offsets[3]  # 3
-    # arr[offset] = -1  # TODO: doesn't work when arr is a num.ndarray
+    arr[offset] = -1
     arr[offset - 2 : offset] = [-1, -1]
-    assert np.array_equal(arr, [4, -1, -1, 1, 0])
+    assert np.array_equal(arr, [4, -1, -1, -1, 0])
 
 
 if __name__ == "__main__":