fix isnan and isfinite errors

Author: nayyirahsan

pykokkos/__init__.py

@@ -17,6 +17,8 @@
     set_device_id,
 )
 
+from pykokkos.lib.ufuncs import _isnan as isnan, _isfinite as isfinite
+
 from pykokkos.lib.info import iinfo, finfo
 from pykokkos.lib.create import zeros, zeros_like, ones, ones_like, full, full_like
 from pykokkos.lib.manipulate import reshape, ravel, expand_dims

pykokkos/lib/ufuncs.py

@@ -145,3 +145,57 @@ def _equal(view1, view2, profiler_name: Optional[str] = None):
         view2=view2,
     )
     return out
+
+def _isnan(view, profiler_name: Optional[str] = None):
+    dtype = view.dtype
+    ndims = len(view.shape)
+    if ndims > 2:
+        raise NotImplementedError("isnan() ufunc only supports up to 2D views")
+    out = pk.View([*view.shape], dtype=pk.bool)
+    if view.shape == ():
+        tid = 1
+    else:
+        tid = view.shape[0]
+    if view.ndim == 0:
+        new_view = pk.View([1], dtype=view.dtype)
+        new_view[0] = view
+        view = new_view
+    _ufunc_kernel_dispatcher(
+        profiler_name=profiler_name,
+        tid=tid,
+        dtype=dtype,
+        ndims=ndims,
+        op="isnan",
+        sub_dispatcher=pk.parallel_for,
+        out=out,
+        view=view,
+    )
+    return out
+
+def _isfinite(view, profiler_name: Optional[str] = None):
+    dtype = view.dtype
+    ndims = len(view.shape)
+    if ndims > 2:
+        raise NotImplementedError("isfinite() ufunc only supports up to 2D views")
+    if view.size == 0:
+        out = pk.View(view.shape, dtype=pk.bool)
+        return out
+    out = pk.View([*view.shape], dtype=pk.bool)
+    if view.shape == ():
+        new_view = pk.View([1], dtype=dtype)
+        new_view[:] = view
+        view = new_view
+        tid = 1
+    else:
+        tid = view.shape[0]
+    _ufunc_kernel_dispatcher(
+        profiler_name=profiler_name,
+        tid=tid,
+        dtype=dtype,
+        ndims=ndims,
+        op="isfinite",
+        sub_dispatcher=pk.parallel_for,
+        out=out,
+        view=view,
+    )
+    return out