diff --git a/pykokkos/lib/ufuncs.py b/pykokkos/lib/ufuncs.py index 3efbb21d..cb1d44a1 100644 --- a/pykokkos/lib/ufuncs.py +++ b/pykokkos/lib/ufuncs.py @@ -11,6 +11,18 @@ def reciprocal_impl_1d_float(tid: int, view: pk.View1D[pk.float]): view[tid] = 1 / view[tid] # type: ignore +@pk.workunit +def reciprocal_impl_2d_double(tid: int, view: pk.View2D[pk.double]): + for i in range(view.extent(1)): # type: ignore + view[tid][i] = 1 / view[tid][i] # type: ignore + + +@pk.workunit +def reciprocal_impl_2d_float(tid: int, view: pk.View2D[pk.float]): + for i in range(view.extent(1)): # type: ignore + view[tid][i] = 1 / view[tid][i] # type: ignore + + def reciprocal(view): """ Return the reciprocal of the argument, element-wise. @@ -33,10 +45,14 @@ def reciprocal(view): """ # see gh-29 for some discussion of the dispatching # awkwardness used here - if str(view.dtype) == "DataType.double": + if str(view.dtype) == "DataType.double" and len(view.shape) == 1: pk.parallel_for(view.shape[0], reciprocal_impl_1d_double, view=view) - elif str(view.dtype) == "DataType.float": + elif str(view.dtype) == "DataType.float" and len(view.shape) == 1: pk.parallel_for(view.shape[0], reciprocal_impl_1d_float, view=view) + elif str(view.dtype) == "DataType.float" and len(view.shape) == 2: + pk.parallel_for(view.shape[0], reciprocal_impl_2d_float, view=view) + elif str(view.dtype) == "DataType.double" and len(view.shape) == 2: + pk.parallel_for(view.shape[0], reciprocal_impl_2d_double, view=view) # NOTE: pretty awkward to both return the view # and operate on it in place; the former is closer # to NumPy semantics diff --git a/tests/test_ufuncs.py b/tests/test_ufuncs.py index cd76a21f..f5efe70e 100644 --- a/tests/test_ufuncs.py +++ b/tests/test_ufuncs.py @@ -1,6 +1,7 @@ import pykokkos as pk import numpy as np +from numpy.random import default_rng from numpy.testing import assert_allclose import pytest @@ -349,3 +350,24 @@ def test_caching(): view[:] = np.arange(10, dtype=np.float32) actual = pk.reciprocal(view=view) assert_allclose(actual, expected) + + +@pytest.mark.parametrize("pk_ufunc, numpy_ufunc", [ + (pk.reciprocal, np.reciprocal), +]) +@pytest.mark.parametrize("pk_dtype, numpy_dtype", [ + (pk.double, np.float64), + (pk.float, np.float32), +]) +def test_2d_exposed_ufuncs_vs_numpy(pk_ufunc, + numpy_ufunc, + pk_dtype, + numpy_dtype): + rng = default_rng(123) + in_arr = rng.random((5, 5)).astype(numpy_dtype) + expected = numpy_ufunc(in_arr) + + view: pk.View2d = pk.View([5, 5], pk_dtype) + view[:] = in_arr + actual = pk_ufunc(view=view) + assert_allclose(actual, expected)