pyopencl 2025.2.5__cp311-cp311-macosx_11_0_arm64.whl → 2025.2.6__cp311-cp311-macosx_11_0_arm64.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of pyopencl might be problematic. Click here for more details.
- pyopencl/_cl.cpython-311-darwin.so +0 -0
- pyopencl/_cl.pyi +10 -7
- pyopencl/_monkeypatch.py +40 -5
- pyopencl/algorithm.py +1 -1
- pyopencl/array.py +209 -124
- pyopencl/cache.py +1 -1
- pyopencl/characterize/__init__.py +2 -4
- pyopencl/clmath.py +0 -1
- pyopencl/cltypes.py +42 -27
- pyopencl/elementwise.py +223 -113
- pyopencl/scan.py +30 -25
- pyopencl/tools.py +326 -211
- {pyopencl-2025.2.5.dist-info → pyopencl-2025.2.6.dist-info}/METADATA +3 -4
- {pyopencl-2025.2.5.dist-info → pyopencl-2025.2.6.dist-info}/RECORD +16 -16
- {pyopencl-2025.2.5.dist-info → pyopencl-2025.2.6.dist-info}/WHEEL +0 -0
- {pyopencl-2025.2.5.dist-info → pyopencl-2025.2.6.dist-info}/licenses/LICENSE +0 -0
pyopencl/array.py
CHANGED
|
@@ -1,7 +1,5 @@
|
|
|
1
1
|
"""CL device arrays."""
|
|
2
2
|
|
|
3
|
-
# NOTE: for elwise_kernel_runner which adds keyword arguments
|
|
4
|
-
# pylint:disable=unexpected-keyword-arg
|
|
5
3
|
from __future__ import annotations
|
|
6
4
|
|
|
7
5
|
|
|
@@ -41,12 +39,13 @@ from typing import (
|
|
|
41
39
|
Concatenate,
|
|
42
40
|
Literal,
|
|
43
41
|
ParamSpec,
|
|
42
|
+
TypeAlias,
|
|
44
43
|
cast,
|
|
45
44
|
)
|
|
46
45
|
from warnings import warn
|
|
47
46
|
|
|
48
47
|
import numpy as np
|
|
49
|
-
from typing_extensions import Self, override
|
|
48
|
+
from typing_extensions import Self, TypeIs, override
|
|
50
49
|
|
|
51
50
|
import pyopencl as cl
|
|
52
51
|
import pyopencl.elementwise as elementwise
|
|
@@ -213,7 +212,7 @@ def _splay(
|
|
|
213
212
|
group_count = max_groups
|
|
214
213
|
work_items_per_group = max_work_items
|
|
215
214
|
|
|
216
|
-
# print("n
|
|
215
|
+
# print(f"n:{n} gc:{group_count} wipg:{work_items_per_group}")
|
|
217
216
|
return (group_count*work_items_per_group,), (work_items_per_group,)
|
|
218
217
|
|
|
219
218
|
|
|
@@ -246,14 +245,13 @@ def elwise_kernel_runner(
|
|
|
246
245
|
assert queue is not None
|
|
247
246
|
|
|
248
247
|
knl = kernel_getter(out, *args, **kwargs)
|
|
249
|
-
|
|
248
|
+
gs, ls = out._get_sizes(queue, knl.get_work_group_info(
|
|
250
249
|
cl.kernel_work_group_info.WORK_GROUP_SIZE,
|
|
251
250
|
queue.device))
|
|
252
|
-
gs, ls = out._get_sizes(queue, work_group_info)
|
|
253
251
|
|
|
254
252
|
knl_args = (out, *args, out.size)
|
|
255
253
|
if ARRAY_KERNEL_EXEC_HOOK is not None:
|
|
256
|
-
return ARRAY_KERNEL_EXEC_HOOK(
|
|
254
|
+
return ARRAY_KERNEL_EXEC_HOOK(
|
|
257
255
|
knl, queue, gs, ls, *knl_args, wait_for=wait_for)
|
|
258
256
|
else:
|
|
259
257
|
return knl(queue, gs, ls, *knl_args, wait_for=wait_for)
|
|
@@ -289,6 +287,12 @@ _ARRAY_GET_SIZES_CACHE: \
|
|
|
289
287
|
_BOOL_DTYPE = np.dtype(np.int8)
|
|
290
288
|
_NOT_PRESENT = object()
|
|
291
289
|
|
|
290
|
+
ScalarLike: TypeAlias = int | float | complex | np.number[Any]
|
|
291
|
+
|
|
292
|
+
|
|
293
|
+
def _is_scalar(s: object) -> TypeIs[ScalarLike]:
|
|
294
|
+
return isinstance(s, SCALAR_CLASSES)
|
|
295
|
+
|
|
292
296
|
|
|
293
297
|
class Array:
|
|
294
298
|
"""A :class:`numpy.ndarray` work-alike that stores its data and performs
|
|
@@ -478,16 +482,6 @@ class Array:
|
|
|
478
482
|
|
|
479
483
|
__array_priority__: ClassVar[int] = 100
|
|
480
484
|
|
|
481
|
-
queue: cl.CommandQueue | None
|
|
482
|
-
shape: tuple[int, ...]
|
|
483
|
-
dtype: np.dtype[Any]
|
|
484
|
-
strides: tuple[int, ...]
|
|
485
|
-
events: list[cl.Event]
|
|
486
|
-
nbytes: int
|
|
487
|
-
size: int
|
|
488
|
-
allocator: Allocator | None
|
|
489
|
-
base_data: cl.MemoryObjectHolder | cl.SVMPointer | None
|
|
490
|
-
|
|
491
485
|
def __init__(
|
|
492
486
|
self,
|
|
493
487
|
cq: cl.Context | cl.CommandQueue | None,
|
|
@@ -568,7 +562,7 @@ class Array:
|
|
|
568
562
|
dtype = np.dtype(dtype)
|
|
569
563
|
|
|
570
564
|
try:
|
|
571
|
-
shape = tuple(shape)
|
|
565
|
+
shape = tuple(shape)
|
|
572
566
|
except TypeError as err:
|
|
573
567
|
if not isinstance(shape, (int, np.integer)):
|
|
574
568
|
raise TypeError(
|
|
@@ -587,7 +581,7 @@ class Array:
|
|
|
587
581
|
raise ValueError(f"negative dimensions are not allowed: {shape}")
|
|
588
582
|
if np.any([np.array([s]).dtype.kind not in ["u", "i"] for s in shape]):
|
|
589
583
|
raise ValueError(
|
|
590
|
-
"
|
|
584
|
+
f"invalid shape {shape} (all dimensions must be integers)")
|
|
591
585
|
size = np.prod(shape_array, dtype=np.uint64).item()
|
|
592
586
|
|
|
593
587
|
if strides is None:
|
|
@@ -624,34 +618,37 @@ class Array:
|
|
|
624
618
|
if alloc_nbytes < 0:
|
|
625
619
|
raise ValueError("cannot allocate CL buffer with negative size")
|
|
626
620
|
|
|
627
|
-
|
|
628
|
-
self.shape = shape
|
|
629
|
-
self.dtype = dtype
|
|
630
|
-
self.strides = strides
|
|
631
|
-
self.events = [] if events is None else events
|
|
632
|
-
self.nbytes = alloc_nbytes
|
|
633
|
-
self.size = size
|
|
634
|
-
self.allocator = allocator
|
|
635
|
-
|
|
621
|
+
base_data = None
|
|
636
622
|
if data is None:
|
|
637
623
|
if alloc_nbytes == 0:
|
|
638
|
-
|
|
624
|
+
base_data = None
|
|
639
625
|
|
|
640
626
|
else:
|
|
641
|
-
if
|
|
627
|
+
if allocator is None:
|
|
642
628
|
if context is None and queue is not None:
|
|
643
629
|
context = queue.context
|
|
644
630
|
|
|
645
|
-
|
|
646
|
-
|
|
631
|
+
assert context is not None
|
|
632
|
+
base_data = cl.Buffer(
|
|
633
|
+
context, cl.mem_flags.READ_WRITE, alloc_nbytes)
|
|
647
634
|
else:
|
|
648
|
-
|
|
635
|
+
base_data = allocator(alloc_nbytes)
|
|
649
636
|
else:
|
|
650
|
-
|
|
651
|
-
|
|
652
|
-
self.
|
|
653
|
-
self.context = context
|
|
654
|
-
self.
|
|
637
|
+
base_data = data
|
|
638
|
+
|
|
639
|
+
self.queue: cl.CommandQueue | None = queue
|
|
640
|
+
self.context: cl.Context | None = context
|
|
641
|
+
self.shape: tuple[int, ...] = shape
|
|
642
|
+
self.dtype: np.dtype[Any] = dtype
|
|
643
|
+
self.strides: tuple[int, ...] = strides
|
|
644
|
+
self.events: list[cl.Event] = [] if events is None else events
|
|
645
|
+
self.nbytes: int = alloc_nbytes
|
|
646
|
+
self.size: int = size
|
|
647
|
+
self.allocator: Allocator | None = allocator
|
|
648
|
+
self.base_data: cl.MemoryObjectHolder | cl.SVMPointer | None = base_data
|
|
649
|
+
self.offset: int = offset
|
|
650
|
+
|
|
651
|
+
self._flags: _ArrayFlags | None = _flags
|
|
655
652
|
|
|
656
653
|
if __debug__:
|
|
657
654
|
if queue is not None and isinstance(
|
|
@@ -664,11 +661,11 @@ class Array:
|
|
|
664
661
|
InconsistentOpenCLQueueWarning, stacklevel=2)
|
|
665
662
|
|
|
666
663
|
@property
|
|
667
|
-
def ndim(self):
|
|
664
|
+
def ndim(self) -> int:
|
|
668
665
|
return len(self.shape)
|
|
669
666
|
|
|
670
667
|
@property
|
|
671
|
-
def data(self):
|
|
668
|
+
def data(self) -> cl.MemoryObjectHolder | cl.SVMPointer | None:
|
|
672
669
|
if self.offset:
|
|
673
670
|
raise ArrayHasOffsetError()
|
|
674
671
|
else:
|
|
@@ -679,6 +676,7 @@ class Array:
|
|
|
679
676
|
f = self._flags
|
|
680
677
|
if f is None:
|
|
681
678
|
self._flags = f = _ArrayFlags(self)
|
|
679
|
+
|
|
682
680
|
return f
|
|
683
681
|
|
|
684
682
|
def _new_with_changes(self,
|
|
@@ -899,14 +897,16 @@ class Array:
|
|
|
899
897
|
|
|
900
898
|
return result
|
|
901
899
|
|
|
902
|
-
|
|
900
|
+
@override
|
|
901
|
+
def __str__(self) -> str:
|
|
903
902
|
if self.queue is None:
|
|
904
903
|
return (f"<cl.{type(self).__name__} {self.shape} of {self.dtype} "
|
|
905
904
|
"without queue, call with_queue()>")
|
|
906
905
|
|
|
907
906
|
return str(self.get())
|
|
908
907
|
|
|
909
|
-
|
|
908
|
+
@override
|
|
909
|
+
def __repr__(self) -> str:
|
|
910
910
|
if self.queue is None:
|
|
911
911
|
return (f"<cl.{type(self).__name__} {self.shape} of {self.dtype} "
|
|
912
912
|
f"at {id(self):x} without queue, call with_queue()>")
|
|
@@ -922,49 +922,69 @@ class Array:
|
|
|
922
922
|
|
|
923
923
|
return result
|
|
924
924
|
|
|
925
|
-
def safely_stringify_for_pudb(self):
|
|
925
|
+
def safely_stringify_for_pudb(self) -> str:
|
|
926
926
|
return f"cl.{type(self).__name__} {self.dtype} {self.shape}"
|
|
927
927
|
|
|
928
|
-
|
|
928
|
+
@override
|
|
929
|
+
def __hash__(self) -> int:
|
|
929
930
|
raise TypeError("pyopencl arrays are not hashable.")
|
|
930
931
|
|
|
931
932
|
# {{{ kernel invocation wrappers
|
|
932
933
|
|
|
933
934
|
@staticmethod
|
|
934
935
|
@elwise_kernel_runner
|
|
935
|
-
def _axpbyz(out
|
|
936
|
-
|
|
936
|
+
def _axpbyz(out: Array,
|
|
937
|
+
a: ScalarLike,
|
|
938
|
+
x: Array,
|
|
939
|
+
b: ScalarLike,
|
|
940
|
+
y: Array,
|
|
941
|
+
queue: cl.CommandQueue | None = None) -> cl.Kernel:
|
|
942
|
+
"""Compute ``out = a*x + b*y``,
|
|
937
943
|
where *other* is an array."""
|
|
938
|
-
|
|
939
|
-
|
|
944
|
+
x_shape = x.shape
|
|
945
|
+
y_shape = y.shape
|
|
940
946
|
out_shape = out.shape
|
|
941
|
-
|
|
942
|
-
|
|
943
|
-
|
|
947
|
+
|
|
948
|
+
assert out.context is not None
|
|
949
|
+
assert (x_shape == y_shape == out_shape
|
|
950
|
+
or (x_shape == () and y_shape == out_shape)
|
|
951
|
+
or (y_shape == () and x_shape == out_shape))
|
|
952
|
+
|
|
944
953
|
return elementwise.get_axpbyz_kernel(
|
|
945
|
-
out.context,
|
|
946
|
-
x_is_scalar=(
|
|
947
|
-
y_is_scalar=(
|
|
954
|
+
out.context, x.dtype, y.dtype, out.dtype,
|
|
955
|
+
x_is_scalar=(x_shape == ()),
|
|
956
|
+
y_is_scalar=(y_shape == ()))
|
|
948
957
|
|
|
949
958
|
@staticmethod
|
|
950
959
|
@elwise_kernel_runner
|
|
951
|
-
def _axpbz(out
|
|
960
|
+
def _axpbz(out: Array,
|
|
961
|
+
a: ScalarLike,
|
|
962
|
+
x: Array,
|
|
963
|
+
b: ScalarLike,
|
|
964
|
+
queue: cl.CommandQueue | None = None) -> cl.Kernel:
|
|
952
965
|
"""Compute ``z = a * x + b``, where *b* is a scalar."""
|
|
953
|
-
a = np.array(a)
|
|
954
|
-
b = np.array(b)
|
|
955
966
|
assert out.shape == x.shape
|
|
956
|
-
|
|
957
|
-
|
|
967
|
+
assert out.context is not None
|
|
968
|
+
|
|
969
|
+
return elementwise.get_axpbz_kernel(
|
|
970
|
+
out.context,
|
|
971
|
+
np.array(a).dtype, x.dtype, np.array(b).dtype, out.dtype)
|
|
958
972
|
|
|
959
973
|
@staticmethod
|
|
960
974
|
@elwise_kernel_runner
|
|
961
|
-
def _elwise_multiply(out
|
|
975
|
+
def _elwise_multiply(out: Array,
|
|
976
|
+
a: Array,
|
|
977
|
+
b: Array,
|
|
978
|
+
queue: cl.CommandQueue | None = None) -> cl.Kernel:
|
|
962
979
|
a_shape = a.shape
|
|
963
980
|
b_shape = b.shape
|
|
964
981
|
out_shape = out.shape
|
|
982
|
+
|
|
965
983
|
assert (a_shape == b_shape == out_shape
|
|
966
984
|
or (a_shape == () and b_shape == out_shape)
|
|
967
985
|
or (b_shape == () and a_shape == out_shape))
|
|
986
|
+
assert a.context is not None
|
|
987
|
+
|
|
968
988
|
return elementwise.get_multiply_kernel(
|
|
969
989
|
a.context, a.dtype, b.dtype, out.dtype,
|
|
970
990
|
x_is_scalar=(a_shape == ()),
|
|
@@ -973,19 +993,27 @@ class Array:
|
|
|
973
993
|
|
|
974
994
|
@staticmethod
|
|
975
995
|
@elwise_kernel_runner
|
|
976
|
-
def _rdiv_scalar(out
|
|
977
|
-
|
|
996
|
+
def _rdiv_scalar(out: Array,
|
|
997
|
+
ary: Array,
|
|
998
|
+
other: ScalarLike,
|
|
999
|
+
queue: cl.CommandQueue | None = None) -> cl.Kernel:
|
|
1000
|
+
assert out.context is not None
|
|
978
1001
|
assert out.shape == ary.shape
|
|
1002
|
+
|
|
979
1003
|
return elementwise.get_rdivide_elwise_kernel(
|
|
980
|
-
out.context, ary.dtype, other.dtype, out.dtype)
|
|
1004
|
+
out.context, ary.dtype, np.array(other).dtype, out.dtype)
|
|
981
1005
|
|
|
982
1006
|
@staticmethod
|
|
983
1007
|
@elwise_kernel_runner
|
|
984
|
-
def _div(out
|
|
1008
|
+
def _div(out: Array,
|
|
1009
|
+
self: Array,
|
|
1010
|
+
other: Array,
|
|
1011
|
+
queue: cl.CommandQueue | None = None) -> cl.Kernel:
|
|
985
1012
|
"""Divides an array by another array."""
|
|
986
1013
|
assert (self.shape == other.shape == out.shape
|
|
987
1014
|
or (self.shape == () and other.shape == out.shape)
|
|
988
1015
|
or (other.shape == () and self.shape == out.shape))
|
|
1016
|
+
assert self.context is not None
|
|
989
1017
|
|
|
990
1018
|
return elementwise.get_divide_kernel(self.context,
|
|
991
1019
|
self.dtype, other.dtype, out.dtype,
|
|
@@ -994,15 +1022,18 @@ class Array:
|
|
|
994
1022
|
|
|
995
1023
|
@staticmethod
|
|
996
1024
|
@elwise_kernel_runner
|
|
997
|
-
def _fill(result, scalar):
|
|
1025
|
+
def _fill(result: Array, scalar: ScalarLike) -> cl.Kernel:
|
|
1026
|
+
assert result.context is not None
|
|
998
1027
|
return elementwise.get_fill_kernel(result.context, result.dtype)
|
|
999
1028
|
|
|
1000
1029
|
@staticmethod
|
|
1001
1030
|
@elwise_kernel_runner
|
|
1002
|
-
def _abs(result, arg):
|
|
1031
|
+
def _abs(result: Array, arg: Array) -> cl.Kernel:
|
|
1032
|
+
assert arg.context is not None
|
|
1033
|
+
|
|
1003
1034
|
if arg.dtype.kind == "c":
|
|
1004
1035
|
from pyopencl.elementwise import complex_dtype_to_name
|
|
1005
|
-
fname = "
|
|
1036
|
+
fname = f"{complex_dtype_to_name(arg.dtype)}_abs"
|
|
1006
1037
|
elif arg.dtype.kind == "f":
|
|
1007
1038
|
fname = "fabs"
|
|
1008
1039
|
elif arg.dtype.kind in ["u", "i"]:
|
|
@@ -1015,63 +1046,83 @@ class Array:
|
|
|
1015
1046
|
|
|
1016
1047
|
@staticmethod
|
|
1017
1048
|
@elwise_kernel_runner
|
|
1018
|
-
def _real(result, arg):
|
|
1049
|
+
def _real(result: Array, arg: Array) -> cl.Kernel:
|
|
1019
1050
|
from pyopencl.elementwise import complex_dtype_to_name
|
|
1020
|
-
|
|
1051
|
+
|
|
1052
|
+
assert arg.context is not None
|
|
1053
|
+
fname = f"{complex_dtype_to_name(arg.dtype)}_real"
|
|
1054
|
+
|
|
1021
1055
|
return elementwise.get_unary_func_kernel(
|
|
1022
1056
|
arg.context, fname, arg.dtype, out_dtype=result.dtype)
|
|
1023
1057
|
|
|
1024
1058
|
@staticmethod
|
|
1025
1059
|
@elwise_kernel_runner
|
|
1026
|
-
def _imag(result, arg):
|
|
1060
|
+
def _imag(result: Array, arg: Array) -> cl.Kernel:
|
|
1027
1061
|
from pyopencl.elementwise import complex_dtype_to_name
|
|
1028
|
-
|
|
1062
|
+
|
|
1063
|
+
assert arg.context is not None
|
|
1064
|
+
fname = f"{complex_dtype_to_name(arg.dtype)}_imag"
|
|
1065
|
+
|
|
1029
1066
|
return elementwise.get_unary_func_kernel(
|
|
1030
1067
|
arg.context, fname, arg.dtype, out_dtype=result.dtype)
|
|
1031
1068
|
|
|
1032
1069
|
@staticmethod
|
|
1033
1070
|
@elwise_kernel_runner
|
|
1034
|
-
def _conj(result, arg):
|
|
1071
|
+
def _conj(result: Array, arg: Array) -> cl.Kernel:
|
|
1035
1072
|
from pyopencl.elementwise import complex_dtype_to_name
|
|
1036
|
-
|
|
1073
|
+
|
|
1074
|
+
assert arg.context is not None
|
|
1075
|
+
fname = f"{complex_dtype_to_name(arg.dtype)}_conj"
|
|
1076
|
+
|
|
1037
1077
|
return elementwise.get_unary_func_kernel(
|
|
1038
1078
|
arg.context, fname, arg.dtype, out_dtype=result.dtype)
|
|
1039
1079
|
|
|
1040
1080
|
@staticmethod
|
|
1041
1081
|
@elwise_kernel_runner
|
|
1042
|
-
def _pow_scalar(result
|
|
1043
|
-
|
|
1082
|
+
def _pow_scalar(result: Array,
|
|
1083
|
+
ary: Array,
|
|
1084
|
+
exponent: ScalarLike) -> cl.Kernel:
|
|
1085
|
+
assert result.context is not None
|
|
1044
1086
|
return elementwise.get_pow_kernel(result.context,
|
|
1045
|
-
ary.dtype, exponent.dtype, result.dtype,
|
|
1087
|
+
ary.dtype, np.array(exponent).dtype, result.dtype,
|
|
1046
1088
|
is_base_array=True, is_exp_array=False)
|
|
1047
1089
|
|
|
1048
1090
|
@staticmethod
|
|
1049
1091
|
@elwise_kernel_runner
|
|
1050
|
-
def _rpow_scalar(result
|
|
1051
|
-
|
|
1092
|
+
def _rpow_scalar(result: Array,
|
|
1093
|
+
base: ScalarLike,
|
|
1094
|
+
exponent: Array) -> cl.Kernel:
|
|
1095
|
+
assert result.context is not None
|
|
1052
1096
|
return elementwise.get_pow_kernel(result.context,
|
|
1053
|
-
base.dtype, exponent.dtype, result.dtype,
|
|
1097
|
+
np.array(base).dtype, exponent.dtype, result.dtype,
|
|
1054
1098
|
is_base_array=False, is_exp_array=True)
|
|
1055
1099
|
|
|
1056
1100
|
@staticmethod
|
|
1057
1101
|
@elwise_kernel_runner
|
|
1058
|
-
def _pow_array(result
|
|
1102
|
+
def _pow_array(result: Array,
|
|
1103
|
+
base: Array,
|
|
1104
|
+
exponent: Array) -> cl.Kernel:
|
|
1105
|
+
assert result.context is not None
|
|
1059
1106
|
return elementwise.get_pow_kernel(
|
|
1060
1107
|
result.context, base.dtype, exponent.dtype, result.dtype,
|
|
1061
1108
|
is_base_array=True, is_exp_array=True)
|
|
1062
1109
|
|
|
1063
1110
|
@staticmethod
|
|
1064
1111
|
@elwise_kernel_runner
|
|
1065
|
-
def _reverse(result, ary):
|
|
1112
|
+
def _reverse(result: Array, ary: Array) -> cl.Kernel:
|
|
1113
|
+
assert result.context is not None
|
|
1066
1114
|
return elementwise.get_reverse_kernel(result.context, ary.dtype)
|
|
1067
1115
|
|
|
1068
1116
|
@staticmethod
|
|
1069
1117
|
@elwise_kernel_runner
|
|
1070
|
-
def _copy(dest, src):
|
|
1118
|
+
def _copy(dest: Array, src: Array) -> cl.Kernel:
|
|
1119
|
+
assert dest.context is not None
|
|
1071
1120
|
return elementwise.get_copy_kernel(
|
|
1072
1121
|
dest.context, dest.dtype, src.dtype)
|
|
1073
1122
|
|
|
1074
|
-
def _new_like_me(self,
|
|
1123
|
+
def _new_like_me(self,
|
|
1124
|
+
dtype: DTypeLike = None,
|
|
1125
|
+
queue: cl.CommandQueue | None = None) -> Self:
|
|
1075
1126
|
if dtype is None:
|
|
1076
1127
|
dtype = self.dtype
|
|
1077
1128
|
strides = self.strides
|
|
@@ -1095,20 +1146,26 @@ class Array:
|
|
|
1095
1146
|
|
|
1096
1147
|
@staticmethod
|
|
1097
1148
|
@elwise_kernel_runner
|
|
1098
|
-
def _scalar_binop(out, a, b
|
|
1149
|
+
def _scalar_binop(out: Array, a: Array, b: ScalarLike, op: str,
|
|
1150
|
+
queue: cl.CommandQueue | None = None) -> cl.Kernel:
|
|
1151
|
+
assert out.context is not None
|
|
1099
1152
|
return elementwise.get_array_scalar_binop_kernel(
|
|
1100
1153
|
out.context, op, out.dtype, a.dtype,
|
|
1101
1154
|
np.array(b).dtype)
|
|
1102
1155
|
|
|
1103
1156
|
@staticmethod
|
|
1104
1157
|
@elwise_kernel_runner
|
|
1105
|
-
def _array_binop(out, a, b
|
|
1158
|
+
def _array_binop(out: Array, a: Array, b: Array, op: str,
|
|
1159
|
+
queue: cl.CommandQueue | None = None) -> cl.Kernel:
|
|
1106
1160
|
a_shape = a.shape
|
|
1107
1161
|
b_shape = b.shape
|
|
1108
1162
|
out_shape = out.shape
|
|
1163
|
+
|
|
1109
1164
|
assert (a_shape == b_shape == out_shape
|
|
1110
1165
|
or (a_shape == () and b_shape == out_shape)
|
|
1111
1166
|
or (b_shape == () and a_shape == out_shape))
|
|
1167
|
+
assert out.context is not None
|
|
1168
|
+
|
|
1112
1169
|
return elementwise.get_array_binop_kernel(
|
|
1113
1170
|
out.context, op, out.dtype, a.dtype, b.dtype,
|
|
1114
1171
|
a_is_scalar=(a_shape == ()),
|
|
@@ -1116,9 +1173,12 @@ class Array:
|
|
|
1116
1173
|
|
|
1117
1174
|
@staticmethod
|
|
1118
1175
|
@elwise_kernel_runner
|
|
1119
|
-
def _unop(out, a
|
|
1176
|
+
def _unop(out: Array, a: Array, op: str,
|
|
1177
|
+
queue: cl.CommandQueue | None = None) -> cl.Kernel:
|
|
1178
|
+
assert out.context is not None
|
|
1120
1179
|
if out.shape != a.shape:
|
|
1121
1180
|
raise ValueError("shapes of arguments do not match")
|
|
1181
|
+
|
|
1122
1182
|
return elementwise.get_unop_kernel(
|
|
1123
1183
|
out.context, op, a.dtype, out.dtype)
|
|
1124
1184
|
|
|
@@ -1590,7 +1650,7 @@ class Array:
|
|
|
1590
1650
|
result.add_event(self._reverse(result, self))
|
|
1591
1651
|
return result
|
|
1592
1652
|
|
|
1593
|
-
def astype(self, dtype, queue: cl.CommandQueue | None = None):
|
|
1653
|
+
def astype(self, dtype: DTypeLike, queue: cl.CommandQueue | None = None):
|
|
1594
1654
|
"""Return a copy of *self*, cast to *dtype*."""
|
|
1595
1655
|
if dtype == self.dtype:
|
|
1596
1656
|
return self.copy()
|
|
@@ -1636,15 +1696,26 @@ class Array:
|
|
|
1636
1696
|
|
|
1637
1697
|
@staticmethod
|
|
1638
1698
|
@elwise_kernel_runner
|
|
1639
|
-
def _scalar_comparison(out
|
|
1699
|
+
def _scalar_comparison(out: Array,
|
|
1700
|
+
a: Array,
|
|
1701
|
+
b: ScalarLike,
|
|
1702
|
+
op: str,
|
|
1703
|
+
queue: cl.CommandQueue | None = None) -> cl.Kernel:
|
|
1704
|
+
assert out.context is not None
|
|
1640
1705
|
return elementwise.get_array_scalar_comparison_kernel(
|
|
1641
1706
|
out.context, op, a.dtype)
|
|
1642
1707
|
|
|
1643
1708
|
@staticmethod
|
|
1644
1709
|
@elwise_kernel_runner
|
|
1645
|
-
def _array_comparison(out
|
|
1710
|
+
def _array_comparison(out: Array,
|
|
1711
|
+
a: Array,
|
|
1712
|
+
b: Array,
|
|
1713
|
+
op: str,
|
|
1714
|
+
queue: cl.CommandQueue | None = None) -> cl.Kernel:
|
|
1715
|
+
assert out.context is not None
|
|
1646
1716
|
if a.shape != b.shape:
|
|
1647
1717
|
raise ValueError("shapes of comparison arguments do not match")
|
|
1718
|
+
|
|
1648
1719
|
return elementwise.get_array_comparison_kernel(
|
|
1649
1720
|
out.context, op, a.dtype, b.dtype)
|
|
1650
1721
|
|
|
@@ -1810,8 +1881,7 @@ class Array:
|
|
|
1810
1881
|
|
|
1811
1882
|
order = kwargs.pop("order", "C")
|
|
1812
1883
|
if kwargs:
|
|
1813
|
-
raise TypeError("unexpected keyword arguments:
|
|
1814
|
-
% list(kwargs.keys()))
|
|
1884
|
+
raise TypeError(f"unexpected keyword arguments: {list(kwargs)}")
|
|
1815
1885
|
|
|
1816
1886
|
if order not in "CF":
|
|
1817
1887
|
raise ValueError("order must be either 'C' or 'F'")
|
|
@@ -2123,8 +2193,7 @@ class Array:
|
|
|
2123
2193
|
index_entry += array_shape
|
|
2124
2194
|
|
|
2125
2195
|
if not (0 <= index_entry < array_shape):
|
|
2126
|
-
raise IndexError(
|
|
2127
|
-
"subindex in axis %d out of range" % index_axis)
|
|
2196
|
+
raise IndexError(f"subindex in axis {index_axis} out of range")
|
|
2128
2197
|
|
|
2129
2198
|
new_offset += self.strides[array_axis]*index_entry
|
|
2130
2199
|
|
|
@@ -2154,7 +2223,7 @@ class Array:
|
|
|
2154
2223
|
index_axis += 1
|
|
2155
2224
|
|
|
2156
2225
|
else:
|
|
2157
|
-
raise IndexError("invalid subindex in axis
|
|
2226
|
+
raise IndexError(f"invalid subindex in axis {index_axis}")
|
|
2158
2227
|
|
|
2159
2228
|
while array_axis < len(self.shape):
|
|
2160
2229
|
new_shape.append(self.shape[array_axis])
|
|
@@ -2656,11 +2725,15 @@ def multi_put(
|
|
|
2656
2725
|
return []
|
|
2657
2726
|
|
|
2658
2727
|
from pytools import single_valued
|
|
2728
|
+
|
|
2659
2729
|
a_dtype = single_valued(a.dtype for a in arrays)
|
|
2660
2730
|
a_allocator = arrays[0].allocator
|
|
2731
|
+
|
|
2661
2732
|
context = dest_indices.context
|
|
2662
2733
|
queue = queue or dest_indices.queue
|
|
2663
2734
|
assert queue is not None
|
|
2735
|
+
assert context is not None
|
|
2736
|
+
|
|
2664
2737
|
if wait_for is None:
|
|
2665
2738
|
wait_for = []
|
|
2666
2739
|
wait_for = [*wait_for, *dest_indices.events]
|
|
@@ -2764,7 +2837,6 @@ def concatenate(arrays, axis=0, queue: cl.CommandQueue | None = None, allocator=
|
|
|
2764
2837
|
raise ValueError(
|
|
2765
2838
|
f"{i_ary}-th array has residual not matching other arrays")
|
|
2766
2839
|
|
|
2767
|
-
# pylint: disable=unsupported-assignment-operation
|
|
2768
2840
|
shape[axis] += ary.shape[axis]
|
|
2769
2841
|
|
|
2770
2842
|
# }}}
|
|
@@ -2960,18 +3032,18 @@ def _if_positive(result, criterion, then_, else_):
|
|
|
2960
3032
|
|
|
2961
3033
|
|
|
2962
3034
|
def if_positive(
|
|
2963
|
-
criterion,
|
|
2964
|
-
then_,
|
|
2965
|
-
else_,
|
|
2966
|
-
out=None,
|
|
3035
|
+
criterion: Array | ScalarLike,
|
|
3036
|
+
then_: Array | ScalarLike,
|
|
3037
|
+
else_: Array | ScalarLike,
|
|
3038
|
+
out: Array | None = None,
|
|
2967
3039
|
queue: cl.CommandQueue | None = None):
|
|
2968
3040
|
"""Return an array like *then_*, which, for the element at index *i*,
|
|
2969
3041
|
contains *then_[i]* if *criterion[i]>0*, else *else_[i]*.
|
|
2970
3042
|
"""
|
|
2971
3043
|
|
|
2972
|
-
is_then_scalar =
|
|
2973
|
-
is_else_scalar =
|
|
2974
|
-
if
|
|
3044
|
+
is_then_scalar = _is_scalar(then_)
|
|
3045
|
+
is_else_scalar = _is_scalar(else_)
|
|
3046
|
+
if _is_scalar(criterion) and is_then_scalar and is_else_scalar:
|
|
2975
3047
|
result = np.where(criterion, then_, else_)
|
|
2976
3048
|
|
|
2977
3049
|
if out is not None:
|
|
@@ -2981,53 +3053,60 @@ def if_positive(
|
|
|
2981
3053
|
return result
|
|
2982
3054
|
|
|
2983
3055
|
if is_then_scalar:
|
|
2984
|
-
|
|
3056
|
+
then_ary = np.array(then_)
|
|
3057
|
+
else:
|
|
3058
|
+
then_ary = then_
|
|
3059
|
+
|
|
3060
|
+
assert not _is_scalar(criterion)
|
|
2985
3061
|
|
|
2986
3062
|
if is_else_scalar:
|
|
2987
|
-
|
|
3063
|
+
else_ary = np.array(else_)
|
|
3064
|
+
else:
|
|
3065
|
+
else_ary = else_
|
|
2988
3066
|
|
|
2989
|
-
if
|
|
3067
|
+
if then_ary.dtype != else_ary.dtype:
|
|
2990
3068
|
raise ValueError(
|
|
2991
|
-
f"dtypes do not match:
|
|
2992
|
-
f"
|
|
3069
|
+
f"dtypes do not match: then_ary is '{then_ary.dtype}' and "
|
|
3070
|
+
f"else_ary is '{else_ary.dtype}'")
|
|
2993
3071
|
|
|
2994
|
-
if
|
|
3072
|
+
if then_ary.shape == () and else_ary.shape == ():
|
|
2995
3073
|
pass
|
|
2996
|
-
elif
|
|
2997
|
-
if not (criterion.shape ==
|
|
3074
|
+
elif then_ary.shape != () and else_ary.shape != ():
|
|
3075
|
+
if not (criterion.shape == then_ary.shape == else_ary.shape):
|
|
2998
3076
|
raise ValueError(
|
|
2999
3077
|
f"shapes do not match: 'criterion' has shape {criterion.shape}"
|
|
3000
|
-
f", '
|
|
3001
|
-
f"{
|
|
3002
|
-
elif
|
|
3003
|
-
if criterion.shape !=
|
|
3078
|
+
f", 'then_ary' has shape {then_ary.shape} and 'else_ary' has shape "
|
|
3079
|
+
f"{else_ary.shape}")
|
|
3080
|
+
elif then_ary.shape == ():
|
|
3081
|
+
if criterion.shape != else_ary.shape:
|
|
3004
3082
|
raise ValueError(
|
|
3005
3083
|
f"shapes do not match: 'criterion' has shape {criterion.shape}"
|
|
3006
|
-
f" and '
|
|
3007
|
-
elif
|
|
3008
|
-
if criterion.shape !=
|
|
3084
|
+
f" and 'else_ary' has shape {else_ary.shape}")
|
|
3085
|
+
elif else_ary.shape == ():
|
|
3086
|
+
if criterion.shape != then_ary.shape:
|
|
3009
3087
|
raise ValueError(
|
|
3010
3088
|
f"shapes do not match: 'criterion' has shape {criterion.shape}"
|
|
3011
|
-
f" and '
|
|
3089
|
+
f" and 'then_ary' has shape {then_ary.shape}")
|
|
3012
3090
|
else:
|
|
3013
3091
|
raise AssertionError()
|
|
3014
3092
|
|
|
3015
3093
|
if out is None:
|
|
3016
|
-
if
|
|
3094
|
+
if then_ary.shape != ():
|
|
3095
|
+
assert isinstance(then_ary, Array)
|
|
3017
3096
|
out = empty_like(
|
|
3018
|
-
|
|
3097
|
+
then_ary, criterion.queue, allocator=criterion.allocator)
|
|
3019
3098
|
else:
|
|
3020
3099
|
# Use same strides as criterion
|
|
3021
3100
|
cr_byte_strides = np.array(criterion.strides, dtype=np.int64)
|
|
3022
3101
|
cr_item_strides = cr_byte_strides // criterion.dtype.itemsize
|
|
3023
|
-
out_strides = tuple(cr_item_strides*
|
|
3102
|
+
out_strides = tuple(cr_item_strides*then_ary.dtype.itemsize)
|
|
3024
3103
|
|
|
3025
3104
|
out = type(criterion)(
|
|
3026
|
-
criterion.queue, criterion.shape,
|
|
3105
|
+
criterion.queue, criterion.shape, then_ary.dtype,
|
|
3027
3106
|
allocator=criterion.allocator,
|
|
3028
3107
|
strides=out_strides)
|
|
3029
3108
|
|
|
3030
|
-
event1 = _if_positive(out, criterion,
|
|
3109
|
+
event1 = _if_positive(out, criterion, then_ary, else_ary, queue=queue)
|
|
3031
3110
|
out.add_event(event1)
|
|
3032
3111
|
|
|
3033
3112
|
return out
|
|
@@ -3105,7 +3184,11 @@ def minimum(a, b, out=None, queue: cl.CommandQueue | None = None):
|
|
|
3105
3184
|
|
|
3106
3185
|
# {{{ logical ops
|
|
3107
3186
|
|
|
3108
|
-
def _logical_op(x1
|
|
3187
|
+
def _logical_op(x1: Array | ScalarLike,
|
|
3188
|
+
x2: Array | ScalarLike,
|
|
3189
|
+
out: Array | None,
|
|
3190
|
+
operator: str,
|
|
3191
|
+
queue: cl.CommandQueue | None = None) -> Array:
|
|
3109
3192
|
# NOTE: Copied from pycuda.gpuarray
|
|
3110
3193
|
assert operator in ["&&", "||"]
|
|
3111
3194
|
|
|
@@ -3129,6 +3212,7 @@ def _logical_op(x1, x2, out, operator, queue: cl.CommandQueue | None = None):
|
|
|
3129
3212
|
|
|
3130
3213
|
out = out or ary_arg._new_like_me(dtype=np.int8)
|
|
3131
3214
|
|
|
3215
|
+
assert queue is not None
|
|
3132
3216
|
assert out.shape == ary_arg.shape and out.dtype == np.int8
|
|
3133
3217
|
|
|
3134
3218
|
knl = elementwise.get_array_scalar_binop_kernel(
|
|
@@ -3153,6 +3237,7 @@ def _logical_op(x1, x2, out, operator, queue: cl.CommandQueue | None = None):
|
|
|
3153
3237
|
out = empty(queue, allocator=allocator,
|
|
3154
3238
|
shape=x1.shape, dtype=np.int8)
|
|
3155
3239
|
|
|
3240
|
+
assert queue is not None
|
|
3156
3241
|
assert out.shape == x1.shape and out.dtype == np.int8
|
|
3157
3242
|
|
|
3158
3243
|
knl = elementwise.get_array_binop_kernel(
|