Support np.select operation

[manopapad]

Python bits are done and C++ stubs are in place, but need to be filled in. Probably a good idea to copy from choose operation. Passing to @ipdemes to take a look.

@bryevdv I am seeing mypy errors on my local machine, that don't appear related to my changes. Possibly a mypy version bump is causing this, can you confirm that you're seeing them too, maybe suggest fixes?

mypy errors

cunumeric/deferred.py: note: In member "storage" of class "DeferredArray":
cunumeric/deferred.py:276:20: error: Incompatible return value type (got
"RegionField | Future", expected "Future | tuple[Region, FieldID]")
[return-value]
                return storage
                       ^~~~~~~
cunumeric/deferred.py:278:20: error: Incompatible return value type (got
"tuple[Region | Any, int | Any]", expected "Future | tuple[Region, FieldID]")
[return-value]
                return (storage.region, storage.field.field_id)
                       ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
cunumeric/deferred.py:278:21: error: Item "Future" of "RegionField | Future"
has no attribute "region"  [union-attr]
                return (storage.region, storage.field.field_id)
                        ^~~~~~~~~~~~~~
cunumeric/deferred.py:278:37: error: Item "Future" of "RegionField | Future"
has no attribute "field"  [union-attr]
                return (storage.region, storage.field.field_id)
                                        ^~~~~~~~~~~~~
cunumeric/deferred.py: note: In member "get_scalar_array" of class "DeferredArray":
cunumeric/deferred.py:405:15: error: Item "RegionField" of
"RegionField | Future" has no attribute "get_buffer"  [union-attr]
            buf = self.base.storage.get_buffer(self.dtype.itemsize)
                  ^~~~~~~~~~~~~~~~~~~~~~~~~~~~
cunumeric/deferred.py: note: In member "_create_indexing_array" of class "DeferredArray":
cunumeric/deferred.py:826:21: error: Unsupported operand types for + ("generic"
and "int")  [operator]
                        k += store.shape[dim + shift]
                        ^
cunumeric/deferred.py:826:21: error: Unsupported operand types for +
("memoryview" and "int")  [operator]
                        k += store.shape[dim + shift]
                        ^
cunumeric/deferred.py:826:21: note: Left operand is of type "generic | bool | int | float | complex | str | bytes | memoryview"
cunumeric/deferred.py:826:26: error: Unsupported operand types for + ("str" and
"int")  [operator]
                        k += store.shape[dim + shift]
                             ^~~~~~~~~~~~~~~~~~~~~~~~
cunumeric/deferred.py:826:26: error: Unsupported operand types for + ("bytes"
and "int")  [operator]
                        k += store.shape[dim + shift]
                             ^~~~~~~~~~~~~~~~~~~~~~~~
cunumeric/deferred.py: note: In member "_get_view" of class "DeferredArray":
cunumeric/deferred.py:903:21: error: Unsupported operand types for + ("generic"
and "int")  [operator]
                        k += store.shape[dim + shift]
                        ^
cunumeric/deferred.py:903:21: error: Unsupported operand types for +
("memoryview" and "int")  [operator]
                        k += store.shape[dim + shift]
                        ^
cunumeric/deferred.py:903:21: note: Left operand is of type "generic | bool | int | float | complex | str | bytes | memoryview"
cunumeric/deferred.py:903:26: error: Unsupported operand types for + ("str" and
"int")  [operator]
                        k += store.shape[dim + shift]
                             ^~~~~~~~~~~~~~~~~~~~~~~~
cunumeric/deferred.py:903:26: error: Unsupported operand types for + ("bytes"
and "int")  [operator]
                        k += store.shape[dim + shift]
                             ^~~~~~~~~~~~~~~~~~~~~~~~
cunumeric/deferred.py: note: In member "swapaxes" of class "DeferredArray":
cunumeric/deferred.py:1332:38: error: Argument 1 to "transpose" of "Store" has
incompatible type "list[int]"; expected "tuple[int, ...]"  [arg-type]
            result = self.base.transpose(dims)
                                         ^~~~
cunumeric/deferred.py:1333:18: error: Incompatible types in assignment
(expression has type "DeferredArray", variable has type "Store")  [assignment]
            result = DeferredArray(self.runtime, result)
                     ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
cunumeric/deferred.py:1335:16: error: Incompatible return value type (got
"Store", expected "DeferredArray")  [return-value]
            return result
                   ^~~~~~
cunumeric/deferred.py: note: In member "choose" of class "DeferredArray":
cunumeric/deferred.py:1741:38: error: Argument 1 to "_broadcast" of
"DeferredArray" has incompatible type "Shape"; expected "tuple[int, ...]"
[arg-type]
            index = index_arr._broadcast(out_arr.shape)
                                         ^~~~~~~~~~~~~
cunumeric/deferred.py:1742:39: error: Argument 1 to "_broadcast" of
"DeferredArray" has incompatible type "Shape"; expected "tuple[int, ...]"
[arg-type]
            ch_tuple = tuple(c._broadcast(out_arr.shape) for c in ch_def)
                                          ^~~~~~~~~~~~~
cunumeric/deferred.py: note: In member "transpose" of class "DeferredArray":
cunumeric/deferred.py:1988:38: error: Argument 1 to "transpose" of "Store" has
incompatible type "tuple[int, ...] | list[int] | None"; expected
"tuple[int, ...]"  [arg-type]
            result = self.base.transpose(axes)
                                         ^~~~
cunumeric/deferred.py:1989:18: error: Incompatible types in assignment
(expression has type "DeferredArray", variable has type "Store")  [assignment]
            result = DeferredArray(self.runtime, result)
                     ^~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
cunumeric/deferred.py:1990:16: error: Incompatible return value type (got
"Store", expected "DeferredArray")  [return-value]
            return result
                   ^~~~~~

[bryevdv]

@manopapad Below is a diff that fixes the mypy issues for me. I am not sure the proper fix for the k += errors since np.isscalar doesn't narrow sufficiently, it seems like stronger assumptions are needed there. The Union[int, FieldID]] seems to mirror what Legate does in many places and offers a fix without coordinating across repos. These unions seem a bit clunky, though.

Edit: FYI I was also just speculating about handling the None axes case this way

computed_axes = tuple(axes) if axes is not None else ()

since transpose does not expect None

diff --git a/cunumeric/config.py b/cunumeric/config.py
index 5020f5a4..f01cbb6b 100644
--- a/cunumeric/config.py
+++ b/cunumeric/config.py
@@ -364,7 +364,7 @@ class CuNumericOpCode(IntEnum):
     SCAN_GLOBAL = _cunumeric.CUNUMERIC_SCAN_GLOBAL
     SCAN_LOCAL = _cunumeric.CUNUMERIC_SCAN_LOCAL
     SEARCHSORTED = _cunumeric.CUNUMERIC_SEARCHSORTED
-    SELECT = _cunumeric.SELECT
+    SELECT = _cunumeric.CUNUMERIC_SELECT
     SOLVE = _cunumeric.CUNUMERIC_SOLVE
     SORT = _cunumeric.CUNUMERIC_SORT
     SYRK = _cunumeric.CUNUMERIC_SYRK
diff --git a/cunumeric/deferred.py b/cunumeric/deferred.py
index fd3844f5..597b899f 100644
--- a/cunumeric/deferred.py
+++ b/cunumeric/deferred.py
@@ -36,6 +36,7 @@ from typing import (
 import legate.core.types as ty
 import numpy as np
 from legate.core import Annotation, Future, ReductionOp, Store
+from legate.core.store import RegionField
 from legate.core.utils import OrderedSet
 from numpy.core.numeric import (  # type: ignore [attr-defined]
     normalize_axis_tuple,
@@ -270,11 +271,13 @@ class DeferredArray(NumPyThunk):
         return f"DeferredArray(base: {self.base})"
 
     @property
-    def storage(self) -> Union[Future, tuple[Region, FieldID]]:
+    def storage(self) -> Union[Future, tuple[Region, Union[int, FieldID]]]:
         storage = self.base.storage
         if self.base.kind == Future:
+            assert isinstance(storage, Future)
             return storage
         else:
+            assert isinstance(storage, RegionField)
             return (storage.region, storage.field.field_id)
 
     @property
@@ -402,6 +405,7 @@ class DeferredArray(NumPyThunk):
 
     def get_scalar_array(self) -> npt.NDArray[Any]:
         assert self.scalar
+        assert isinstance(self.base.storage, Future)
         buf = self.base.storage.get_buffer(self.dtype.itemsize)
         result = np.frombuffer(buf, dtype=self.dtype, count=1)
         return result.reshape(())
@@ -770,10 +774,11 @@ class DeferredArray(NumPyThunk):
 
         store = self.base
         rhs = self
+        computed_key: tuple[Any, ...]
         if isinstance(key, NumPyThunk):
-            key = (key,)
+            computed_key = (key,)
         assert isinstance(key, tuple)
-        key = self._unpack_ellipsis(key, self.ndim)
+        computed_key = self._unpack_ellipsis(key, self.ndim)
 
         # the index where the first index_array is passed to the [] operator
         start_index = -1
@@ -788,7 +793,7 @@ class DeferredArray(NumPyThunk):
         tuple_of_arrays: tuple[Any, ...] = ()
 
         # First, we need to check if transpose is needed
-        for dim, k in enumerate(key):
+        for dim, k in enumerate(computed_key):
             if np.isscalar(k) or isinstance(k, NumPyThunk):
                 if start_index == -1:
                     start_index = dim
@@ -813,17 +818,19 @@ class DeferredArray(NumPyThunk):
             )
             transpose_indices += post_indices
             post_indices = tuple(
-                i for i in range(len(key)) if i not in key_transpose_indices
+                i
+                for i in range(len(computed_key))
+                if i not in key_transpose_indices
             )
             key_transpose_indices += post_indices
             store = store.transpose(transpose_indices)
-            key = tuple(key[i] for i in key_transpose_indices)
+            key = tuple(computed_key[i] for i in key_transpose_indices)
 
         shift = 0
-        for dim, k in enumerate(key):
+        for dim, k in enumerate(computed_key):
             if np.isscalar(k):
                 if k < 0:  # type: ignore [operator]
-                    k += store.shape[dim + shift]
+                    k += store.shape[dim + shift]  # type: ignore [operator]
                 store = store.project(dim + shift, k)
                 shift -= 1
             elif k is np.newaxis:
@@ -831,7 +838,7 @@ class DeferredArray(NumPyThunk):
             elif isinstance(k, slice):
                 k, store = self._slice_store(k, store, dim + shift)
             elif isinstance(k, NumPyThunk):
-                if not isinstance(key, DeferredArray):
+                if not isinstance(computed_key, DeferredArray):
                     k = self.runtime.to_deferred_array(k)
                 if k.dtype == bool:
                     for i in range(k.ndim):
@@ -900,7 +907,7 @@ class DeferredArray(NumPyThunk):
                 k, store = self._slice_store(k, store, dim + shift)
             elif np.isscalar(k):
                 if k < 0:  # type: ignore [operator]
-                    k += store.shape[dim + shift]
+                    k += store.shape[dim + shift]  # type: ignore [operator]
                 store = store.project(dim + shift, k)
                 shift -= 1
             else:
@@ -1329,10 +1336,9 @@ class DeferredArray(NumPyThunk):
         dims = list(range(self.ndim))
         dims[axis1], dims[axis2] = dims[axis2], dims[axis1]
 
-        result = self.base.transpose(dims)
-        result = DeferredArray(self.runtime, result)
+        result = self.base.transpose(tuple(dims))
 
-        return result
+        return DeferredArray(self.runtime, result)
 
     # Convert the source array to the destination array
     @auto_convert("rhs")
@@ -1738,8 +1744,8 @@ class DeferredArray(NumPyThunk):
 
         out_arr = self.base
         # broadcast input array and all choices arrays to the same shape
-        index = index_arr._broadcast(out_arr.shape)
-        ch_tuple = tuple(c._broadcast(out_arr.shape) for c in ch_def)
+        index = index_arr._broadcast(out_arr.shape.extents)
+        ch_tuple = tuple(c._broadcast(out_arr.shape.extents) for c in ch_def)
 
         task = self.context.create_auto_task(CuNumericOpCode.CHOOSE)
         task.add_output(out_arr)
@@ -1985,9 +1991,9 @@ class DeferredArray(NumPyThunk):
     def transpose(
         self, axes: Union[None, tuple[int, ...], list[int]]
     ) -> DeferredArray:
-        result = self.base.transpose(axes)
-        result = DeferredArray(self.runtime, result)
-        return result
+        computed_axes = tuple(axes) if axes is not None else ()
+        result = self.base.transpose(computed_axes)
+        return DeferredArray(self.runtime, result)
 
     @auto_convert("rhs")
     def trilu(self, rhs: Any, k: int, lower: bool) -> None:
diff --git a/cunumeric/eager.py b/cunumeric/eager.py
index 1e913739..56a056b6 100644
--- a/cunumeric/eager.py
+++ b/cunumeric/eager.py
@@ -235,7 +235,7 @@ class EagerArray(NumPyThunk):
         self.escaped = False
 
     @property
-    def storage(self) -> Union[Future, tuple[Region, FieldID]]:
+    def storage(self) -> Union[Future, tuple[Region, Union[int, FieldID]]]:
         if self.deferred is None:
             self.to_deferred_array()
 
diff --git a/cunumeric/thunk.py b/cunumeric/thunk.py
index 0a831bce..68aafb6c 100644
--- a/cunumeric/thunk.py
+++ b/cunumeric/thunk.py
@@ -74,7 +74,7 @@ class NumPyThunk(ABC):
     # Abstract methods
 
     @abstractproperty
-    def storage(self) -> Union[Future, tuple[Region, FieldID]]:
+    def storage(self) -> Union[Future, tuple[Region, Union[int, FieldID]]]:
         """Return the Legion storage primitive for this NumPy thunk"""
         ...

[ipdemes]

@bryevdv : thank you for the fix! I am taking over the work on finishing np.select .
@manopapad: should I continue to use your branch or should I create a local clone?

[manopapad]

I would suggest keeping the same branch for expedience. You should be able to push to it I think

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[manopapad]

/ok to test

[manopapad]

OK, I think I tormented you enough. I think one file was miscommited, but otherwise this is ready to go.

Note that you'll need to approve the PR (because I opened it I can't approve it myself) before you can merge it.

[manopapad]

Was this file added by mistake?

[ipdemes]

yes, sorry