MAINT: PR 57 revisions

* more changes to improve behavior with
array API test suite

pykokkos/__init__.py

@@ -19,6 +19,8 @@
                                  log1p,
                                  sqrt)
 from pykokkos.lib.info import iinfo, finfo
+from pykokkos.lib.create import zeros
+from pykokkos.lib.util import all, any
 
 runtime_singleton.runtime = Runtime()
 defaults: Optional[CompilationDefaults] = runtime_singleton.runtime.compiler.read_defaults()

pykokkos/interface/data_types.py

@@ -30,41 +30,42 @@ class DataTypeClass:
     pass
 
 
-int16 = DataType.int16
+class int16(DataTypeClass):
+    value = kokkos.int16
 
+class int32(DataTypeClass):
+    value = kokkos.int32
 
-int32 = DataType.int32
-
-
-int64 = DataType.int64
-
+class int64(DataTypeClass):
+    value = kokkos.int64
 
 class uint16(DataTypeClass):
-    pass
+    value = kokkos.uint16
 
 
 class uint32(DataTypeClass):
-    pass
+    value = kokkos.int32
 
 
 class uint64(DataTypeClass):
-    pass
-
+    value = kokkos.int64
 
-float = DataType.float
 
+class float(DataTypeClass):
+    value = kokkos.float
 
-double = DataType.double
-double.__name__ = "double" # type: ignore
+class double(DataTypeClass):
+    value = kokkos.double
 
 
 class real(DataTypeClass):
-    pass
+    value = None
 
 class float32(DataTypeClass):
-    pass
+    value = kokkos.float
 
 class float64(DataTypeClass):
-    pass
+    value = kokkos.double
 
-bool = DataType.bool
+class bool(DataTypeClass):
+    value = kokkos.int16

pykokkos/interface/views.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 import ctypes
+import os
 import math
 from enum import Enum
 import sys
@@ -10,6 +11,7 @@
 
 import numpy as np
 
+import pykokkos as pk
 from pykokkos.bindings import kokkos
 import pykokkos.kokkos_manager as km
 
@@ -131,6 +133,8 @@ def __len__(self) -> int:
         :returns: the length of the first dimension
         """
 
+        if len(self.shape) == 0:
+            return 0
         return self.shape[0]
 
     def __iter__(self) -> Iterator:
@@ -240,6 +244,8 @@ def _init_view(
         """
 
         self.shape: List[int] = shape
+        if self.shape == [0]:
+            self.shape = ()
         self.size = math.prod(shape)
         self.dtype: Optional[DataType] = self._get_type(dtype)
         if self.dtype is None:
@@ -259,9 +265,19 @@ def _init_view(
         self.layout: Layout = layout
         self.trait: Trait = trait
 
+        if self.dtype == pk.float:
+            self.dtype = DataType.float
+        elif self.dtype == pk.double:
+            self.dtype = DataType.double
+        elif self.dtype == pk.int32:
+            self.dtype = DataType.int32
+        elif self.dtype == pk.int64:
+            pass
         if trait is trait.Unmanaged:
             self.array = kokkos.unmanaged_array(array, dtype=self.dtype.value, space=self.space.value, layout=self.layout.value)
         else:
+            if len(shape) == 0:
+                shape = [1]
             self.array = kokkos.array("", shape, None, None, self.dtype.value, space.value, layout.value, trait.value)
         self.data = np.array(self.array, copy=False)
 
@@ -283,6 +299,9 @@ def _get_type(self, dtype: Union[DataType, type]) -> Optional[DataType]:
             if dtype is real:
                 return DataType[km.get_default_precision().__name__]
 
+            if dtype == DataType.int64:
+                dtype = int64
+
             return dtype
 
         if dtype is int:
@@ -326,10 +345,14 @@ def __init__(self, parent_view: Union[Subview, View], data_slice: Union[slice, T
         self.base_view: View = self._get_base_view(parent_view)
 
         self.data: np.ndarray = parent_view.data[data_slice]
+        self.dtype = parent_view.dtype
         self.array = kokkos.array(
             self.data, dtype=parent_view.dtype.value, space=parent_view.space.value,
             layout=parent_view.layout.value, trait=kokkos.Unmanaged)
         self.shape: List[int] = list(self.data.shape)
+        if self.data.shape == (0,):
+            self.data = np.array([], dtype=self.dtype)
+            self.shape = ()
 
         self.parent_slice: List[Union[int, slice]]
         self.parent_slice = self._create_slice(data_slice)
@@ -373,6 +396,13 @@ def _get_base_view(self, parent_view: Union[Subview, View]) -> View:
 
         return base_view
 
+    def __eq__(self, other):
+        if isinstance(other, View):
+            if len(self.data) == 0 and len(other.data) == 0:
+                return True
+            result_of_eq = self.data == other.data
+            return result_of_eq
+
 def from_numpy(array: np.ndarray, space: Optional[MemorySpace] = None, layout: Optional[Layout] = None) -> ViewType:
     """
     Create a PyKokkos View from a numpy array
@@ -402,6 +432,8 @@ def from_numpy(array: np.ndarray, space: Optional[MemorySpace] = None, layout: O
         dtype = DataType.float # PyKokkos float
     elif np_dtype is np.float64:
         dtype = double
+    elif np_dtype is np.bool_:
+        dtype = int16
     else:
         raise RuntimeError(f"ERROR: unsupported numpy datatype {np_dtype}")
 
@@ -420,8 +452,8 @@ def from_numpy(array: np.ndarray, space: Optional[MemorySpace] = None, layout: O
     # TODO: pykokkos support for 0-D arrays?
     # temporary/terrible hack here for array API testing..
     if array.ndim == 0:
-        ret_list = list((1,))
-        array = [0]
+        ret_list = ()
+        array = np.array(())
     else:
         ret_list = list((array.shape))
 

pykokkos/lib/create.py

@@ -0,0 +1,4 @@
+import pykokkos as pk
+
+def zeros(shape, *, dtype=None, device=None):
+    return pk.View([*shape], dtype=dtype)

pykokkos/lib/info.py

@@ -25,8 +25,13 @@ def iinfo(type_or_arr):
     # values from the NumPy equivalent
     if "int32" in str(type_or_arr):
         return info_type_attrs(bits=32,
-                               min=2147483647,
-                               max=-2147483648)
+                               max=2147483647,
+                               min=-2147483648)
+    elif "int16" in str(type_or_arr):
+        # iinfo(min=-32768, max=32767, dtype=int16)
+        return info_type_attrs(bits=16,
+                               min=-32768,
+                               max=32767)
     elif "int64" in str(type_or_arr):
         return info_type_attrs(bits=64,
                                min=-9223372036854775808,

pykokkos/lib/util.py

@@ -0,0 +1,21 @@
+import pykokkos as pk
+
+import numpy as np
+
+# TODO: add proper implementations rather
+# than wrapping NumPy
+# These are required for the array API:
+# https://data-apis.org/array-api/2021.12/API_specification/utility_functions.html
+
+def all(x, /, *, axis=None, keepdims=False):
+    if x == True:
+        return True
+    elif x == False:
+        return False
+    np_result = np.all(x)
+    ret_val = pk.View(pk.from_numpy(np.all(x)))
+    return ret_val
+
+
+def any(x, /, *, axis=None, keepdims=False):
+    return pk.View(pk.from_numpy(np.any(x)))