vectorize-pr

cunumeric/__init__.py

@@ -36,6 +36,7 @@
 from cunumeric.logic import *
 from cunumeric.window import bartlett, blackman, hamming, hanning, kaiser
 from cunumeric.coverage import clone_module
+from cunumeric.vectorize import vectorize
 
 clone_module(_np, globals(), maybe_convert_to_np_ndarray)
 

cunumeric/array.py

@@ -33,7 +33,7 @@
 
 import legate.core.types as ty
 import numpy as np
-import pyarrow  # type: ignore  [import]
+import pyarrow
 from legate.core import Array
 from numpy.core.multiarray import (  # type: ignore [attr-defined]
     normalize_axis_index,

cunumeric/config.py

@@ -143,9 +143,11 @@ class _CunumericSharedLib:
     CUNUMERIC_CONVERT_NAN_PROD: int
     CUNUMERIC_CONVERT_NAN_SUM: int
     CUNUMERIC_CONVOLVE: int
+    CUNUMERIC_CREATE_CU_KERNEL: int
     CUNUMERIC_DIAG: int
     CUNUMERIC_DOT: int
     CUNUMERIC_EYE: int
+    CUNUMERIC_EVAL_UDF: int
     CUNUMERIC_FFT: int
     CUNUMERIC_FFT_C2C: int
     CUNUMERIC_FFT_C2R: int
@@ -345,9 +347,11 @@ class CuNumericOpCode(IntEnum):
     CONTRACT = _cunumeric.CUNUMERIC_CONTRACT
     CONVERT = _cunumeric.CUNUMERIC_CONVERT
     CONVOLVE = _cunumeric.CUNUMERIC_CONVOLVE
+    CREATE_CU_KERNEL = _cunumeric.CUNUMERIC_CREATE_CU_KERNEL
     DIAG = _cunumeric.CUNUMERIC_DIAG
     DOT = _cunumeric.CUNUMERIC_DOT
     EYE = _cunumeric.CUNUMERIC_EYE
+    EVAL_UDF = _cunumeric.CUNUMERIC_EVAL_UDF
     FFT = _cunumeric.CUNUMERIC_FFT
     FILL = _cunumeric.CUNUMERIC_FILL
     FLIP = _cunumeric.CUNUMERIC_FLIP

cunumeric/utils.py

@@ -22,6 +22,7 @@
 
 import legate.core.types as ty
 import numpy as np
+import pyarrow as pa
 
 from .types import NdShape
 
@@ -43,6 +44,26 @@
     np.float64: ty.float64,
     np.complex64: ty.complex64,
     np.complex128: ty.complex128,
+    complex: ty.complex128,
+}
+
+CUNUMERIC_TYPE_MAP = {
+    bool: ty.bool_,
+    int: ty.int64,
+    float: ty.float64,
+    complex: ty.complex128,
+    pa.bool_: ty.bool_,
+    pa.int8: ty.int8,
+    pa.int16: ty.int16,
+    pa.int32: ty.int32,
+    pa.int64: ty.int64,  # np.int is int
+    pa.uint8: ty.uint8,
+    pa.uint16: ty.uint16,
+    pa.uint32: ty.uint32,
+    pa.uint64: ty.uint64,  # np.uint is np.uint64
+    pa.float16: ty.float16,
+    pa.float32: ty.float32,
+    pa.float64: ty.float64,
 }
 
 
@@ -99,6 +120,12 @@ def is_supported_dtype(dtype: Any) -> bool:
     return dtype.type in SUPPORTED_DTYPES
 
 
+def convert_to_cunumeric_dtype(dtype: Any) -> Any:
+    if dtype in CUNUMERIC_TYPE_MAP:
+        return CUNUMERIC_TYPE_MAP[dtype]
+    raise TypeError("dtype is not supported")
+
+
 def calculate_volume(shape: NdShape) -> int:
     if len(shape) == 0:
         return 0