Advanced indexing

cunumeric/array.py

@@ -767,8 +767,21 @@ def _convert_key(self, key, first=True):
         elif isinstance(key, tuple) and first:
             return tuple(self._convert_key(k, first=False) for k in key)
         else:
-            # Otherwise convert it to a cuNumeric array and get the thunk
-            return convert_to_cunumeric_ndarray(key)._thunk
+            # Otherwise convert it to a cuNumeric array, check types
+            # and get the thunk
+            key = convert_to_cunumeric_ndarray(key)
+            if key.dtype != np.bool and not np.issubdtype(
+                key.dtype, np.integer
+            ):
+                raise TypeError("index arrays should be int or bool type")
+            if key.dtype != np.bool and key.dtype != np.int64:
+                runtime.warn(
+                    "converting index array to int64 type",
+                    category=RuntimeWarning,
+                )
+                key = key.astype(np.int64)
+
+            return key._thunk
 
     @add_boilerplate()
     def __getitem__(self, key):

cunumeric/config.py

@@ -78,6 +78,7 @@ def destroy(self):
 # Match these to CuNumericOpCode in cunumeric_c.h
 @unique
 class CuNumericOpCode(IntEnum):
+    ADVANCED_INDEXING = _cunumeric.CUNUMERIC_ADVANCED_INDEXING
     ARANGE = _cunumeric.CUNUMERIC_ARANGE
     BINARY_OP = _cunumeric.CUNUMERIC_BINARY_OP
     BINARY_RED = _cunumeric.CUNUMERIC_BINARY_RED
@@ -115,6 +116,7 @@ class CuNumericOpCode(IntEnum):
     WHERE = _cunumeric.CUNUMERIC_WHERE
     WINDOW = _cunumeric.CUNUMERIC_WINDOW
     WRITE = _cunumeric.CUNUMERIC_WRITE
+    ZIP = _cunumeric.CUNUMERIC_ZIP
 
 
 # Match these to CuNumericUnaryOpCode in cunumeric_c.h
@@ -252,3 +254,17 @@ class CuNumericTunable(IntEnum):
     NUM_PROCS = _cunumeric.CUNUMERIC_TUNABLE_NUM_PROCS
     MAX_EAGER_VOLUME = _cunumeric.CUNUMERIC_TUNABLE_MAX_EAGER_VOLUME
     HAS_NUMAMEM = _cunumeric.CUNUMERIC_TUNABLE_HAS_NUMAMEM
+
+
+# Match these to CuNumericTypeCodes in cunumeric_c.h
+@unique
+class CuNumericTypeCodes(IntEnum):
+    CUNUMERIC_TYPE_POINT1 = _cunumeric.CUNUMERIC_TYPE_POINT1
+    CUNUMERIC_TYPE_POINT2 = _cunumeric.CUNUMERIC_TYPE_POINT2
+    CUNUMERIC_TYPE_POINT3 = _cunumeric.CUNUMERIC_TYPE_POINT3
+    CUNUMERIC_TYPE_POINT4 = _cunumeric.CUNUMERIC_TYPE_POINT4
+    CUNUMERIC_TYPE_POINT5 = _cunumeric.CUNUMERIC_TYPE_POINT5
+    CUNUMERIC_TYPE_POINT6 = _cunumeric.CUNUMERIC_TYPE_POINT6
+    CUNUMERIC_TYPE_POINT7 = _cunumeric.CUNUMERIC_TYPE_POINT7
+    CUNUMERIC_TYPE_POINT8 = _cunumeric.CUNUMERIC_TYPE_POINT8
+    CUNUMERIC_TYPE_POINT9 = _cunumeric.CUNUMERIC_TYPE_POINT9