pyopencl 2025.2.4__cp311-cp311-win_amd64.whl → 2025.2.6__cp311-cp311-win_amd64.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/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:%d gc:%d wipg:%d" % (n, group_count, work_items_per_group))
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
- work_group_info = cast("int", knl.get_work_group_info(
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( # pylint: disable=not-callable
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) # type: ignore[arg-type]
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
- "Invalid shape %s ; dimensions, must be integer" % (str(shape)))
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
- self.queue = queue
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
- self.base_data = None
624
+ base_data = None
639
625
 
640
626
  else:
641
- if self.allocator is None:
627
+ if allocator is None:
642
628
  if context is None and queue is not None:
643
629
  context = queue.context
644
630
 
645
- self.base_data = cl.Buffer(
646
- context, cl.mem_flags.READ_WRITE, alloc_nbytes)
631
+ assert context is not None
632
+ base_data = cl.Buffer(
633
+ context, cl.mem_flags.READ_WRITE, alloc_nbytes)
647
634
  else:
648
- self.base_data = self.allocator(alloc_nbytes)
635
+ base_data = allocator(alloc_nbytes)
649
636
  else:
650
- self.base_data = data
651
-
652
- self.offset = offset
653
- self.context = context
654
- self._flags = _flags
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
- def __str__(self):
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
- def __repr__(self):
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
- def __hash__(self):
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, afac, a, bfac, b, queue: cl.CommandQueue | None = None):
936
- """Compute ``out = selffac * self + otherfac*other``,
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
- a_shape = a.shape
939
- b_shape = b.shape
944
+ x_shape = x.shape
945
+ y_shape = y.shape
940
946
  out_shape = out.shape
941
- assert (a_shape == b_shape == out_shape
942
- or (a_shape == () and b_shape == out_shape)
943
- or (b_shape == () and a_shape == out_shape))
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, a.dtype, b.dtype, out.dtype,
946
- x_is_scalar=(a_shape == ()),
947
- y_is_scalar=(b_shape == ()))
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, a, x, b, queue: cl.CommandQueue | None = None):
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
- return elementwise.get_axpbz_kernel(out.context,
957
- a.dtype, x.dtype, b.dtype, out.dtype)
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, a, b, queue: cl.CommandQueue | None = None):
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, ary, other, queue: cl.CommandQueue | None = None):
977
- other = np.array(other)
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, self, other, queue: cl.CommandQueue | None = None):
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 = "%s_abs" % complex_dtype_to_name(arg.dtype)
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
- fname = "%s_real" % complex_dtype_to_name(arg.dtype)
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
- fname = "%s_imag" % complex_dtype_to_name(arg.dtype)
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
- fname = "%s_conj" % complex_dtype_to_name(arg.dtype)
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, ary, exponent):
1043
- exponent = np.array(exponent)
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, base, exponent):
1051
- base = np.array(base)
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, base, exponent):
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, dtype=None, queue: cl.CommandQueue | None = None):
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, queue: cl.CommandQueue | None = None, op=None):
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, queue: cl.CommandQueue | None = None, op=None):
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, queue: cl.CommandQueue | None = None, op=None):
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, a, b, queue: cl.CommandQueue | None = None, op=None):
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, a, b, queue: cl.CommandQueue | None = None, op=None):
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: %s"
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 %d" % index_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])
@@ -2380,7 +2449,7 @@ def _arange_knl(result, start, step):
2380
2449
  result.context, result.dtype)
2381
2450
 
2382
2451
 
2383
- def arange(queue, *args, **kwargs):
2452
+ def arange(queue: cl.CommandQueue, *args: Any, **kwargs: Any) -> Array:
2384
2453
  """arange(queue, [start, ] stop [, step], **kwargs)
2385
2454
  Create a :class:`Array` filled with numbers spaced *step* apart,
2386
2455
  starting from *start* and ending at *stop*. If not given, *start*
@@ -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 = isinstance(then_, SCALAR_CLASSES)
2973
- is_else_scalar = isinstance(else_, SCALAR_CLASSES)
2974
- if isinstance(criterion, SCALAR_CLASSES) and is_then_scalar and is_else_scalar:
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
- then_ = np.array(then_)
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
- else_ = np.array(else_)
3063
+ else_ary = np.array(else_)
3064
+ else:
3065
+ else_ary = else_
2988
3066
 
2989
- if then_.dtype != else_.dtype:
3067
+ if then_ary.dtype != else_ary.dtype:
2990
3068
  raise ValueError(
2991
- f"dtypes do not match: then_ is '{then_.dtype}' and "
2992
- f"else_ is '{else_.dtype}'")
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 then_.shape == () and else_.shape == ():
3072
+ if then_ary.shape == () and else_ary.shape == ():
2995
3073
  pass
2996
- elif then_.shape != () and else_.shape != ():
2997
- if not (criterion.shape == then_.shape == else_.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", 'then_' has shape {then_.shape} and 'else_' has shape "
3001
- f"{else_.shape}")
3002
- elif then_.shape == ():
3003
- if criterion.shape != else_.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 'else_' has shape {else_.shape}")
3007
- elif else_.shape == ():
3008
- if criterion.shape != then_.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 'then_' has shape {then_.shape}")
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 then_.shape != ():
3094
+ if then_ary.shape != ():
3095
+ assert isinstance(then_ary, Array)
3017
3096
  out = empty_like(
3018
- then_, criterion.queue, allocator=criterion.allocator)
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*then_.dtype.itemsize)
3102
+ out_strides = tuple(cr_item_strides*then_ary.dtype.itemsize)
3024
3103
 
3025
3104
  out = type(criterion)(
3026
- criterion.queue, criterion.shape, then_.dtype,
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, then_, else_, queue=queue)
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, x2, out, operator, queue: cl.CommandQueue | None = None):
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(