From 0624001e14951464b9215cb4ab4222fa1638f569 Mon Sep 17 00:00:00 2001 From: Manolis Papadakis Date: Thu, 21 Mar 2024 09:39:10 -0700 Subject: [PATCH] Expose is_supported_dtype to the public interface (#150) Also take this opportunity to clean up a naming inconsistency; NumPy types are "dtypes", core types are "types". --- cunumeric/__init__.py | 1 + cunumeric/_array/array.py | 10 +++++----- cunumeric/_thunk/deferred.py | 4 ++-- cunumeric/_utils/array.py | 17 +++++++++++++++-- cunumeric/runtime.py | 8 ++++---- tests/integration/test_argsort.py | 2 +- tests/integration/test_prod.py | 2 +- tests/integration/test_searchsorted.py | 2 +- tests/unit/cunumeric/test_utils_array.py | 6 +++--- 9 files changed, 33 insertions(+), 19 deletions(-) diff --git a/cunumeric/__init__.py b/cunumeric/__init__.py index 3df52a64c..34c3ada13 100644 --- a/cunumeric/__init__.py +++ b/cunumeric/__init__.py @@ -31,6 +31,7 @@ from ._array.util import maybe_convert_to_np_ndarray from ._module import * from ._ufunc import * +from ._utils.array import is_supported_dtype from ._utils.coverage import clone_module clone_module(_np, globals(), maybe_convert_to_np_ndarray) diff --git a/cunumeric/_array/array.py b/cunumeric/_array/array.py index 729717fe2..b138e75e9 100644 --- a/cunumeric/_array/array.py +++ b/cunumeric/_array/array.py @@ -30,7 +30,7 @@ ) from .. import _ufunc -from .._utils.array import calculate_volume, to_core_dtype +from .._utils.array import calculate_volume, to_core_type from .._utils.coverage import FALLBACK_WARNING, clone_class, is_implemented from .._utils.linalg import dot_modes from .._utils.structure import deep_apply @@ -128,7 +128,7 @@ def __init__( for inp in inputs if isinstance(inp, ndarray) ] - core_dtype = to_core_dtype(dtype) + core_dtype = to_core_type(dtype) self._thunk = runtime.create_empty_thunk( sanitized_shape, core_dtype, inputs ) @@ -660,7 +660,7 @@ def __contains__(self, item: Any) -> ndarray: args = (np.array(item, dtype=self.dtype),) if args[0].size != 1: raise ValueError("contains needs scalar item") - core_dtype = to_core_dtype(self.dtype) + core_dtype = to_core_type(self.dtype) return perform_unary_reduction( UnaryRedCode.CONTAINS, self, @@ -1975,7 +1975,7 @@ def clip( return convert_to_cunumeric_ndarray( self.__array__().clip(args[0], args[1]) ) - core_dtype = to_core_dtype(self.dtype) + core_dtype = to_core_type(self.dtype) extra_args = (Scalar(min, core_dtype), Scalar(max, core_dtype)) return perform_unary_op( UnaryOpCode.CLIP, self, out=out, extra_args=extra_args @@ -2971,7 +2971,7 @@ def var( # FIXME(wonchanl): the following code blocks on mu to convert # it to a Scalar object. We need to get rid of this blocking by # allowing the extra arguments to be Legate stores - args=(Scalar(mu.__array__(), to_core_dtype(self.dtype)),), + args=(Scalar(mu.__array__(), to_core_type(self.dtype)),), ) else: # TODO(https://github.com/nv-legate/cunumeric/issues/591) diff --git a/cunumeric/_thunk/deferred.py b/cunumeric/_thunk/deferred.py index 6fe09523a..ddc5d9641 100644 --- a/cunumeric/_thunk/deferred.py +++ b/cunumeric/_thunk/deferred.py @@ -49,7 +49,7 @@ normalize_axis_tuple, ) -from .._utils.array import is_advanced_indexing, to_core_dtype +from .._utils.array import is_advanced_indexing, to_core_type from ..config import ( BinaryOpCode, BitGeneratorDistribution, @@ -1701,7 +1701,7 @@ def select( c_arr = c._broadcast(self.shape) task.add_input(c_arr) task.add_alignment(c_arr, out_arr) - task.add_scalar_arg(default, to_core_dtype(default.dtype)) + task.add_scalar_arg(default, to_core_type(default.dtype)) task.execute() # Create or extract a diagonal from a matrix diff --git a/cunumeric/_utils/array.py b/cunumeric/_utils/array.py index 1ba3c6bfe..705335cd8 100644 --- a/cunumeric/_utils/array.py +++ b/cunumeric/_utils/array.py @@ -40,11 +40,24 @@ } -def is_supported_type(dtype: str | np.dtype[Any]) -> bool: +def is_supported_dtype(dtype: str | np.dtype[Any]) -> bool: + """ + Whether a NumPy dtype is supported by cuNumeric + + Parameters + ---------- + dtype : data-type + The dtype to query + + Returns + ------- + res : bool + True if `dtype` is a supported dtype + """ return np.dtype(dtype) in SUPPORTED_DTYPES -def to_core_dtype(dtype: str | np.dtype[Any]) -> ty.Type: +def to_core_type(dtype: str | np.dtype[Any]) -> ty.Type: core_dtype = SUPPORTED_DTYPES.get(np.dtype(dtype)) if core_dtype is None: raise TypeError(f"cuNumeric does not support dtype={dtype}") diff --git a/cunumeric/runtime.py b/cunumeric/runtime.py index 8d2f630e2..33db34f9a 100644 --- a/cunumeric/runtime.py +++ b/cunumeric/runtime.py @@ -23,7 +23,7 @@ from legate.core import LEGATE_MAX_DIM, Scalar, TaskTarget, get_legate_runtime from legate.settings import settings as legate_settings -from ._utils.array import calculate_volume, is_supported_type, to_core_dtype +from ._utils.array import calculate_volume, is_supported_dtype, to_core_type from ._utils.stack import find_last_user_stacklevel from .config import ( BitGeneratorOperation, @@ -60,7 +60,7 @@ def thunk_from_scalar( from ._thunk.deferred import DeferredArray store = legate_runtime.create_store_from_scalar( - Scalar(bytes, to_core_dtype(dtype)), + Scalar(bytes, to_core_type(dtype)), shape=shape, ) return DeferredArray(store) @@ -377,7 +377,7 @@ def find_or_create_array_thunk( from ._thunk.deferred import DeferredArray assert isinstance(array, np.ndarray) - if not is_supported_type(array.dtype): + if not is_supported_dtype(array.dtype): raise TypeError(f"cuNumeric does not support dtype={array.dtype}") # We have to be really careful here to handle the case of @@ -429,7 +429,7 @@ def find_or_create_array_thunk( # This is not a scalar so make a field. # We won't try to cache these bigger arrays. store = legate_runtime.create_store_from_buffer( - to_core_dtype(array.dtype), + to_core_type(array.dtype), array.shape, array.copy() if transfer == TransferType.MAKE_COPY else array, # This argument should really be called "donate" diff --git a/tests/integration/test_argsort.py b/tests/integration/test_argsort.py index b36de0d16..656e22b80 100644 --- a/tests/integration/test_argsort.py +++ b/tests/integration/test_argsort.py @@ -98,7 +98,7 @@ def test_structured_array_order(self): # if self.deferred is None: # if self.parent is None: # - # > assert self.runtime.is_supported_type(self.array.dtype) + # > assert self.runtime.is_supported_dtype(self.array.dtype) # E # AssertionError # diff --git a/tests/integration/test_prod.py b/tests/integration/test_prod.py index c004c95a3..4a820905f 100644 --- a/tests/integration/test_prod.py +++ b/tests/integration/test_prod.py @@ -222,7 +222,7 @@ def test_dtype_complex(self, dtype): # allclose hits assertion error: # File "/legate/cunumeric/cunumeric/eager.py", line 293, # in to_deferred_array - # assert self.runtime.is_supported_type(self.array.dtype) + # assert self.runtime.is_supported_dtype(self.array.dtype) # AssertionError assert allclose(out_np, out_num) diff --git a/tests/integration/test_searchsorted.py b/tests/integration/test_searchsorted.py index c3d1461af..03860ec97 100644 --- a/tests/integration/test_searchsorted.py +++ b/tests/integration/test_searchsorted.py @@ -83,7 +83,7 @@ def test_val_none(self): # cuNumeric raises AssertionError # if self.deferred is None: # if self.parent is None: - # > assert self.runtime.is_supported_type + # > assert self.runtime.is_supported_dtype # (self.array.dtype) # E AssertionError # cunumeric/cunumeric/eager.py:to_deferred_array() diff --git a/tests/unit/cunumeric/test_utils_array.py b/tests/unit/cunumeric/test_utils_array.py index 27aa59a12..34e124c47 100644 --- a/tests/unit/cunumeric/test_utils_array.py +++ b/tests/unit/cunumeric/test_utils_array.py @@ -75,17 +75,17 @@ class Test_is_supported_dtype: @pytest.mark.parametrize("value", ["foo", 10, 10.2, (), set()]) def test_type_bad(self, value) -> None: with pytest.raises(TypeError): - m.to_core_dtype(value) + m.to_core_type(value) @pytest.mark.parametrize("value", EXPECTED_SUPPORTED_DTYPES) def test_supported(self, value) -> None: - m.to_core_dtype(value) + m.to_core_type(value) # This is just a representative sample, not exhasutive @pytest.mark.parametrize("value", [np.float128, np.datetime64, [], {}]) def test_unsupported(self, value) -> None: with pytest.raises(TypeError): - m.to_core_dtype(value) + m.to_core_type(value) @pytest.mark.parametrize(