Skip to content

privatize pk.array #406

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
nayyirahsan wants to merge 2 commits into
base: main
Choose a base branch
from
Open

privatize pk.array #406

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
709e7ee
privatize pk.array
nayyirahsan Mar 31, 2026
fc64ad7
fix ci error
nayyirahsan Apr 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions benchmarks/dgemm_compare.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,9 +61,9 @@
b_new = b
c_new = c

view_a = pk.array(a_new)
view_b = pk.array(b_new)
view_c = pk.array(c_new)
view_a = pk.asarray(a_new)
view_b = pk.asarray(b_new)
view_c = pk.asarray(c_new)

pk_dgemm_time_sec = timeit.timeit(
"pk_dgemm(alpha, view_a, view_b, beta, view_c)",
Expand Down
4 changes: 2 additions & 2 deletions examples/NaiveBayes/GaussianNB.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -357,8 +357,8 @@ class labels known to the classifier.
Examples
--------
>>> import numpy as np
>>> X = pk.array([[-1, -1], [-2, -1], [-3, -2], [1, 1], [2, 1], [3, 2]])
>>> Y = pk.array([1, 1, 1, 2, 2, 2])
>>> X = pk.asarray([[-1, -1], [-2, -1], [-3, -2], [1, 1], [2, 1], [3, 2]])
>>> Y = pk.asarray([1, 1, 1, 2, 2, 2])
>>> from sklearn.naive_bayes import GaussianNB
>>> clf = GaussianNB()
>>> clf.fit(X, Y)
Expand Down
2 changes: 1 addition & 1 deletion examples/kokkos/inclusive_scan_team_cuda.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ def main():
num_teams = (N + team_size - 1) // team_size

view = cp.zeros(N, dtype=cp.int32)
view_pk = pk.array(view)
view_pk = pk._array(view)
p_init = pk.RangePolicy(pk.ExecutionSpace.Cuda, 0, N)
pk.parallel_for(p_init, init_data, view=view_pk)

Expand Down
6 changes: 3 additions & 3 deletions examples/pykokkos/from_array.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,9 +27,9 @@ def addition_cp(i: int, cp_arr):
print(f"before {cp_arr=}")
print(f"before {list_arr=}")

np_view = pk.array(np_arr)
cp_view = pk.array(cp_arr)
list_view = pk.array(list_arr)
np_view = pk.asarray(np_arr)
cp_view = pk._array(cp_arr)
list_view = pk.asarray(list_arr)

pk.parallel_for(pk.RangePolicy(pk.OpenMP, 0, size), addition_np, np_arr=np_view)
pk.parallel_for(pk.RangePolicy(pk.Cuda, 0, size), addition_cp, cp_arr=cp_view)
Expand Down
4 changes: 2 additions & 2 deletions examples/pykokkos/multi_gpu.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,14 @@ def reduction_cp(i: int, acc: pk.Acc[int], cp_arr):


pk.set_device_id(1)
cp_view_0 = pk.array(cp_arr_1)
cp_view_0 = pk._array(cp_arr_1)
result_0 = pk.parallel_reduce(
pk.RangePolicy(pk.Cuda, 0, size), reduction_cp, cp_arr=cp_view_0
)
print(result_0)

pk.set_device_id(0)
cp_view_1 = pk.array(cp_arr_0)
cp_view_1 = pk._array(cp_arr_0)
result_1 = pk.parallel_reduce(
pk.RangePolicy(pk.Cuda, 0, size), reduction_cp, cp_arr=cp_view_1
)
Expand Down
1 change: 1 addition & 0 deletions pykokkos/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@
from pykokkos.runtime import runtime_singleton
from pykokkos.core import Runtime
from pykokkos.interface import *
from pykokkos.interface.views import _array
from pykokkos.kokkos_manager import (
initialize,
finalize,
Expand Down
2 changes: 2 additions & 0 deletions pykokkos/interface/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,8 @@
ScratchView6D,
ScratchView7D,
ScratchView8D,
_array,
# Transitional export: deprecated, use pk.asarray for user conversions.
array,
asarray,
result_type,
Expand Down
39 changes: 36 additions & 3 deletions pykokkos/interface/views.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
from enum import Enum
import os
import sys
import warnings
from types import ModuleType
from typing import Dict, Generic, Iterator, List, Optional, Tuple, TypeVar, Union

Expand Down Expand Up @@ -945,7 +946,7 @@ def is_array(array) -> bool:
return True


def array(
def _array(
array, space: Optional[MemorySpace] = None, layout: Optional[Layout] = None
) -> ViewType:
"""
Expand All @@ -958,8 +959,11 @@ def array(
"""

# if numpy array, use from_numpy()
if isinstance(array, np.ndarray) or np.isscalar(array):
if isinstance(array, np.ndarray):
return from_numpy(array, space, layout)
# Python scalars do not expose .dtype, so normalize first
if np.isscalar(array):
return from_numpy(np.asarray(array), space, layout)
# test if the input array can duck-type to a numpy-like array
# and run from_array to preprocess the array to numpy
if is_array(array):
Expand All @@ -968,6 +972,25 @@ def array(
return from_numpy(np.asarray(array), space, layout)


def array(
array, space: Optional[MemorySpace] = None, layout: Optional[Layout] = None
) -> ViewType:
"""
Deprecated public compatibility shim for internal array conversion.

Prefer `pk.asarray` for user-level conversions.
"""

warnings.warn(
"pk.array is deprecated and will be removed in a future release. "
"Use pk.asarray for user conversions; pk._array is internal/private.",
DeprecationWarning,
stacklevel=2,
)

return _array(array, space, layout)

Comment on lines +975 to +992

@IvanGrigorik IvanGrigorik May 15, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can remove it completely.
No need to give warnings.


# asarray is required for comformance with the array API:
# https://data-apis.org/array-api/2021.12/API_specification/creation_functions.html#objects-in-api

Expand All @@ -977,7 +1000,17 @@ def asarray(obj, /, *, dtype=None, device=None, copy=None):
# for now, let's cheat and use NumPy asarray() followed
# by pykokkos from_numpy()

if not isinstance(obj, list) and obj in {pk.e, pk.pi, pk.inf, pk.nan}:
is_array_api_constant = False
if np.isscalar(obj):
if obj in (pk.e, pk.pi, pk.inf):
is_array_api_constant = True
else:
try:
is_array_api_constant = bool(np.isnan(obj))
except TypeError:
is_array_api_constant = False

if is_array_api_constant:
Comment on lines +1003 to +1013

@IvanGrigorik IvanGrigorik May 15, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What's the point of it? Why we can't use pk.nan?
And it looks invalid since we can test cupy arrays with np.isnan, which should give errors.

if dtype is None:
dtype = pk.float64
view = pk.View([1], dtype=dtype)
Expand Down
4 changes: 2 additions & 2 deletions pykokkos/lib/util.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,12 +11,12 @@

def all(x, /, *, axis=None, keepdims=False):
np_result = np.all(x)
ret_val = pk.array(np_result)
ret_val = pk._array(np_result)
return ret_val


def any(x, /, *, axis=None, keepdims=False):
return pk.View(pk.array(np.any(x)))
return pk.View(pk._array(np.any(x)))


@pk.workunit
Expand Down
10 changes: 5 additions & 5 deletions tests/test_linalg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -182,12 +182,12 @@ def test_dgemm_shape(shape_a, shape_b):
def test_dgemm_vs_scipy(alpha, a, b, c, beta, expected_c):
# test against expected results from
# scipy.linalg.blas.dgemm
view_a = pk.array(a)
view_b = pk.array(b)
view_a = pk.asarray(a)
view_b = pk.asarray(b)
if c is None:
view_c = None
else:
view_c = pk.array(c)
view_c = pk.asarray(c)
actual_c = dgemm(
alpha=alpha, view_a=view_a, view_b=view_b, view_c=view_c, beta=beta
)
Expand All @@ -196,7 +196,7 @@ def test_dgemm_vs_scipy(alpha, a, b, c, beta, expected_c):

def test_dgemm_input_handling():
alpha = 1.0
view_a = pk.array(np.zeros((4, 3)))
view_b = pk.array([np.array([0, 0, 0], dtype=np.int32)] * 4)
view_a = pk.asarray(np.zeros((4, 3)))
view_b = pk.asarray([np.array([0, 0, 0], dtype=np.int32)] * 4)
with pytest.raises(ValueError, match="Second dimensions"):
dgemm(alpha=alpha, view_a=view_a, view_b=view_b)
2 changes: 1 addition & 1 deletion tests/test_parallelreduce.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,7 +66,7 @@ def test_squaresum_types(series_max, dtype):
np_data = np.arange(series_max, dtype=dtype)
expected = np.sum(np_data**2)

view = pk.array(np_data)
view = pk.asarray(np_data)
policy = pk.RangePolicy(pk.ExecutionSpace.OpenMP, 0, series_max)

if dtype == np.float64:
Expand Down
30 changes: 15 additions & 15 deletions tests/test_ufuncs.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -259,7 +259,7 @@ def test_multi_array_1d_exposed_ufuncs_vs_numpy(
def test_scalar_operations_vs_numpy(pk_ufunc, numpy_ufunc, numpy_dtype):
data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
expected = numpy_ufunc(np.array(data, dtype=numpy_dtype), 1)
actual = pk_ufunc(pk.array(np.array(data, dtype=numpy_dtype)), 1)
actual = pk_ufunc(pk.asarray(np.array(data, dtype=numpy_dtype)), 1)
assert_allclose(actual, expected)


Expand Down Expand Up @@ -423,8 +423,8 @@ def test_np_matmul_2d_1d_vs_numpy(pk_ufunc, numpy_ufunc, numpy_dtype, test_dim):
np2 = rng.random(N2).astype(numpy_dtype)
expected = numpy_ufunc(np1, np2)

view1 = pk.array(np1)
view2 = pk.array(np2)
view1 = pk.asarray(np1)
view2 = pk.asarray(np2)
actual = pk_ufunc(view1, view2)

assert_allclose(actual, expected)
Expand Down Expand Up @@ -454,8 +454,8 @@ def test_np_matmul_1d_2d_vs_numpy(pk_ufunc, numpy_ufunc, numpy_dtype, test_dim):
np2 = rng.random((N2, M2)).astype(numpy_dtype)
expected = numpy_ufunc(np1, np2)

view1 = pk.array(np1)
view2 = pk.array(np2)
view1 = pk.asarray(np1)
view2 = pk.asarray(np2)
actual = pk_ufunc(view1, view2)

assert_allclose(actual, expected)
Expand Down Expand Up @@ -495,8 +495,8 @@ def test_np_matmul_fails(numpy_dtype, test_dim):
np2 = rng.random((N2, M2)).astype(numpy_dtype)

with pytest.raises(RuntimeError) as e_info:
view1 = pk.array(np1)
view2 = pk.array(np2)
view1 = pk.asarray(np1)
view2 = pk.asarray(np2)
pk.np_matmul(view1, view2) # Should fail with 1d x 2d

err_np_matmul = (
Expand Down Expand Up @@ -536,8 +536,8 @@ def test_multi_array_2d_exposed_ufuncs_vs_numpy(pk_ufunc, numpy_ufunc, numpy_dty
np2 = rng.random((N, M)).astype(numpy_dtype)
expected = numpy_ufunc(np1, np2)

view1 = pk.array(np1)
view2 = pk.array(np2)
view1 = pk.asarray(np1)
view2 = pk.asarray(np2)
actual = pk_ufunc(view1, view2)

assert_allclose(actual, expected)
Expand Down Expand Up @@ -592,8 +592,8 @@ def test_broadcast_array_exposed_ufuncs_vs_numpy(

expected = numpy_ufunc(np1, np2)

view1 = pk.array(np1)
view2 = pk.array(np2) if isinstance(np2, np.ndarray) else np2
view1 = pk.asarray(np1)
view2 = pk.asarray(np2) if isinstance(np2, np.ndarray) else np2
actual = pk_ufunc(view1, view2)

assert_allclose(expected, actual)
Expand Down Expand Up @@ -643,8 +643,8 @@ def test_copyto_1d(pk_ufunc, numpy_ufunc, numpy_dtype):
np2 = rng.random((N, M)).astype(numpy_dtype)
numpy_ufunc(np1, np2)

view1 = pk.array(np1)
view2 = pk.array(np2)
view1 = pk.asarray(np1)
view2 = pk.asarray(np2)
pk_ufunc(view1, view2)

assert_allclose(np1, view1)
Expand Down Expand Up @@ -705,8 +705,8 @@ def test_copyto_broadcast_2d(pk_ufunc, numpy_ufunc, numpy_dtype, test_dim):

numpy_ufunc(np1, np2)

view1 = pk.array(np1)
view2 = pk.array(np2) if isinstance(np2, np.ndarray) else np2
view1 = pk.asarray(np1)
view2 = pk.asarray(np2) if isinstance(np2, np.ndarray) else np2
pk_ufunc(view1, view2)

assert_allclose(np1, view1)
Expand Down
49 changes: 42 additions & 7 deletions tests/test_views.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,9 +97,9 @@ def __init__(self, threads: int, i_1: int, i_2: int, i_3: int, i_4: int):
cp_arr = cp.zeros((threads, 2)).astype(np.int32)
list_arr = [np.array([0, 0], dtype=np.int32)] * threads

self.np_view: pk.View2D[int] = pk.array(np_arr)
self.cp_view: pk.View2D[int] = pk.array(cp_arr)
self.list_view: pk.View2D[int] = pk.array(list_arr)
self.np_view: pk.View2D[int] = pk.asarray(np_arr)
self.cp_view: pk.View2D[int] = pk._array(cp_arr)
self.list_view: pk.View2D[int] = pk.asarray(list_arr)
Comment on lines +100 to +102

@IvanGrigorik IvanGrigorik May 15, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The end goal is not to use pykokkos internal data structures. If we will deprecate array but still will use asarray - that's not correct.
We want to have something like self.np_view = np.array(...) and be able to use it like that.
To be 100% flexible, we can have some intermediate data structure (per-file only) like xp that should decide if we want to use cupy or numpy based on current execution space.
But there should not be xp.asarray(numpy_array). It should be xp.array(...).

You can treat xp as a C++ macros like this:

#ifdef __CUDA_ARCH__
    xp = cupy
#else
    xp = numpy
#endif


@pk.workunit
def v1d(self, tid: int) -> None:
Expand Down Expand Up @@ -283,9 +283,9 @@ def test_arrays(self):
cp_arr = cp.zeros((self.threads, 2)).astype(np.int32)
list_arr = [np.array([0, 0], dtype=np.int32)] * self.threads

np_view = pk.array(np_arr)
cp_view = pk.array(cp_arr)
list_view = pk.array(list_arr)
np_view = pk.asarray(np_arr)
cp_view = pk._array(cp_arr)
list_view = pk.asarray(list_arr)

pk.parallel_for(
pk.RangePolicy(pk.OpenMP, 0, self.threads), addition_np, np_arr=np_view
Expand Down Expand Up @@ -376,6 +376,41 @@ def test_asarray_consts_vs_numpy(const, np_dtype, pk_dtype):
assert not "int" in pk_type_string


@pytest.mark.parametrize(
"arr",
[
np.array([1, 2, 3], dtype=np.int32),
[1, 2, 3],
7,
],
)
def test_array_deprecated_alias_still_converts(arr):
with pytest.warns(DeprecationWarning, match="pk.array is deprecated"):
view = pk.array(arr)
assert_allclose(view, np.asarray(arr))


@pytest.mark.parametrize(
"arr",
[
np.array([3, 1, 4], dtype=np.int32),
[3, 1, 4],
],
)
def test_private_array_conversion_matches_asarray(arr):
from_private = pk._array(arr)
from_public = pk.asarray(arr)
assert type(from_private) is type(from_public)
assert_allclose(from_private, from_public)


@pytest.mark.skipif(not HAS_CUDA, reason="CUDA/cupy not available")

@IvanGrigorik IvanGrigorik May 15, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This one is cool. We should add more of this + more CUDA test.

def test_private_array_conversion_accepts_cupy_array():
cp_arr = cp.array([3, 1, 4], dtype=cp.int32)
from_private = pk._array(cp_arr)
assert_allclose(from_private, cp.asnumpy(cp_arr))


@pytest.mark.parametrize(
"pk_dtype, np_dtype",
[
Expand Down Expand Up @@ -416,7 +451,7 @@ def test_result_type_supported(pk_dtype, pk_dtype2, expected_promo):
@pytest.mark.parametrize(
"pk_dtype, pk_dtype2",
[
(pk.array(np.array([0])), pk.uint16),
(pk.asarray(np.array([0])), pk.uint16),
(pk.uint64, pk.int8),
(pk.float32, pk.int64),
],
Expand Down
Loading