gpuarray 0.1.0__tar.gz → 0.2.0__tar.gz

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.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: gpuarray
3
- Version: 0.1.0
3
+ Version: 0.2.0
4
4
  Summary: NumPy-like GPU array library that works on every GPU — NVIDIA, AMD, Intel, Apple Silicon. No CUDA required.
5
5
  Author: Arnav Kewalram
6
6
  License: MIT
@@ -20,6 +20,13 @@ from .api import (
20
20
  relu,
21
21
  sigmoid,
22
22
  tanh,
23
+ where,
24
+ clip,
25
+ argmax,
26
+ argmin,
27
+ concatenate,
28
+ stack,
29
+ random,
23
30
  )
24
31
 
25
32
  __version__ = "0.1.0"
@@ -44,4 +51,11 @@ __all__ = [
44
51
  "relu",
45
52
  "sigmoid",
46
53
  "tanh",
54
+ "where",
55
+ "clip",
56
+ "argmax",
57
+ "argmin",
58
+ "concatenate",
59
+ "stack",
60
+ "random",
47
61
  ]
@@ -0,0 +1,182 @@
1
+ """NumPy-compatible module-level functions."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import Sequence
6
+
7
+ import numpy as np
8
+
9
+ from .core import GPUArray, _from_numpy, matmul as _matmul
10
+ from . import shaders
11
+
12
+
13
+ # ----- creation functions -----
14
+
15
+ def array(data) -> GPUArray:
16
+ """Create a GPUArray from a list, nested list, or numpy array."""
17
+ if isinstance(data, GPUArray):
18
+ return data
19
+ arr = np.asarray(data, dtype=np.float32)
20
+ if arr.ndim == 0:
21
+ arr = arr.reshape((1,))
22
+ return _from_numpy(arr)
23
+
24
+
25
+ def zeros(shape) -> GPUArray:
26
+ if isinstance(shape, int):
27
+ shape = (shape,)
28
+ return _from_numpy(np.zeros(shape, dtype=np.float32))
29
+
30
+
31
+ def ones(shape) -> GPUArray:
32
+ if isinstance(shape, int):
33
+ shape = (shape,)
34
+ return _from_numpy(np.ones(shape, dtype=np.float32))
35
+
36
+
37
+ def arange(start, stop=None, step=1) -> GPUArray:
38
+ if stop is None:
39
+ stop = start
40
+ start = 0
41
+ return _from_numpy(np.arange(start, stop, step, dtype=np.float32))
42
+
43
+
44
+ def linspace(start, stop, num=50) -> GPUArray:
45
+ return _from_numpy(np.linspace(start, stop, num, dtype=np.float32))
46
+
47
+
48
+ # ----- operations -----
49
+
50
+ def dot(a: GPUArray, b: GPUArray):
51
+ return a.dot(b)
52
+
53
+
54
+ def matmul(a: GPUArray, b: GPUArray) -> GPUArray:
55
+ return _matmul(a, b)
56
+
57
+
58
+ def sum(a: GPUArray) -> float:
59
+ return a.sum()
60
+
61
+
62
+ def mean(a: GPUArray) -> float:
63
+ return a.mean()
64
+
65
+
66
+ def exp(a: GPUArray) -> GPUArray:
67
+ return a._run_unary_op(shaders.EXP_SHADER)
68
+
69
+
70
+ def log(a: GPUArray) -> GPUArray:
71
+ return a._run_unary_op(shaders.LOG_SHADER)
72
+
73
+
74
+ def sqrt(a: GPUArray) -> GPUArray:
75
+ return a._run_unary_op(shaders.SQRT_SHADER)
76
+
77
+
78
+ def abs(a: GPUArray) -> GPUArray:
79
+ return a._run_unary_op(shaders.ABS_SHADER)
80
+
81
+
82
+ def max(a: GPUArray) -> float:
83
+ return a.max()
84
+
85
+
86
+ def min(a: GPUArray) -> float:
87
+ return a.min()
88
+
89
+
90
+ def relu(a: GPUArray) -> GPUArray:
91
+ return a._run_unary_op(shaders.RELU_SHADER)
92
+
93
+
94
+ def sigmoid(a: GPUArray) -> GPUArray:
95
+ return a._run_unary_op(shaders.SIGMOID_SHADER)
96
+
97
+
98
+ def tanh(a: GPUArray) -> GPUArray:
99
+ return a._run_unary_op(shaders.TANH_SHADER)
100
+
101
+
102
+ # ----- new operations -----
103
+
104
+ def where(condition: GPUArray, x: GPUArray, y: GPUArray) -> GPUArray:
105
+ """Conditional select: where condition is 1.0, pick x; else pick y."""
106
+ import wgpu
107
+ from .device import get_device
108
+ from .core import _get_pipeline
109
+
110
+ assert condition.size == x.size == y.size
111
+ device = get_device()
112
+ _, pipeline = _get_pipeline(shaders.WHERE_SHADER)
113
+
114
+ out_buf = device.create_buffer(
115
+ size=condition.size * 4,
116
+ usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST,
117
+ )
118
+ bind_group = device.create_bind_group(
119
+ layout=pipeline.get_bind_group_layout(0),
120
+ entries=[
121
+ {"binding": 0, "resource": {"buffer": condition._buffer, "offset": 0, "size": condition.size * 4}},
122
+ {"binding": 1, "resource": {"buffer": x._buffer, "offset": 0, "size": x.size * 4}},
123
+ {"binding": 2, "resource": {"buffer": y._buffer, "offset": 0, "size": y.size * 4}},
124
+ {"binding": 3, "resource": {"buffer": out_buf, "offset": 0, "size": condition.size * 4}},
125
+ ],
126
+ )
127
+ import math
128
+ n_workgroups = math.ceil(condition.size / 256)
129
+ encoder = device.create_command_encoder()
130
+ pass_enc = encoder.begin_compute_pass()
131
+ pass_enc.set_pipeline(pipeline)
132
+ pass_enc.set_bind_group(0, bind_group)
133
+ pass_enc.dispatch_workgroups(n_workgroups)
134
+ pass_enc.end()
135
+ device.queue.submit([encoder.finish()])
136
+ return GPUArray(out_buf, condition._shape)
137
+
138
+
139
+ def clip(a: GPUArray, min_val: float, max_val: float) -> GPUArray:
140
+ """Clamp values to [min_val, max_val]."""
141
+ return a.clip(min_val, max_val)
142
+
143
+
144
+ def argmax(a: GPUArray) -> int:
145
+ """Return index of maximum value."""
146
+ return a.argmax()
147
+
148
+
149
+ def argmin(a: GPUArray) -> int:
150
+ """Return index of minimum value."""
151
+ return a.argmin()
152
+
153
+
154
+ def concatenate(arrays, axis: int = 0) -> GPUArray:
155
+ """Join arrays along an axis (CPU-side, creates new GPU buffer)."""
156
+ np_arrays = [a.to_numpy() for a in arrays]
157
+ result = np.concatenate(np_arrays, axis=axis)
158
+ return _from_numpy(result)
159
+
160
+
161
+ def stack(arrays, axis: int = 0) -> GPUArray:
162
+ """Stack arrays along a new axis (CPU-side, creates new GPU buffer)."""
163
+ np_arrays = [a.to_numpy() for a in arrays]
164
+ result = np.stack(np_arrays, axis=axis)
165
+ return _from_numpy(result)
166
+
167
+
168
+ # ----- random submodule -----
169
+
170
+ class _Random:
171
+ """Random array generation (CPU-side, uploaded to GPU)."""
172
+
173
+ def rand(self, *shape) -> GPUArray:
174
+ """Uniform random values in [0, 1)."""
175
+ return _from_numpy(np.random.rand(*shape).astype(np.float32))
176
+
177
+ def randn(self, *shape) -> GPUArray:
178
+ """Standard normal random values."""
179
+ return _from_numpy(np.random.randn(*shape).astype(np.float32))
180
+
181
+
182
+ random = _Random()
@@ -343,6 +343,105 @@ class GPUArray:
343
343
  raise ValueError(f"Cannot reshape array of size {self.size} into shape {new_shape}")
344
344
  return GPUArray(self._buffer, new_shape)
345
345
 
346
+ # ----- clamp / clip -----
347
+
348
+ def clip(self, min_val: float, max_val: float) -> "GPUArray":
349
+ """Clamp values to [min_val, max_val]."""
350
+ device = get_device()
351
+ _, pipeline = _get_pipeline(shaders.CLAMP_SHADER)
352
+
353
+ out_buf = device.create_buffer(
354
+ size=self.size * 4,
355
+ usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST,
356
+ )
357
+ # Pack two f32 values, padded to 16 bytes
358
+ params_data = struct.pack("ff", min_val, max_val) + b"\x00" * 8
359
+ uniform_buf = device.create_buffer_with_data(
360
+ data=params_data,
361
+ usage=wgpu.BufferUsage.UNIFORM,
362
+ )
363
+ bind_group = device.create_bind_group(
364
+ layout=pipeline.get_bind_group_layout(0),
365
+ entries=[
366
+ {"binding": 0, "resource": {"buffer": self._buffer, "offset": 0, "size": self.size * 4}},
367
+ {"binding": 1, "resource": {"buffer": out_buf, "offset": 0, "size": self.size * 4}},
368
+ {"binding": 2, "resource": {"buffer": uniform_buf, "offset": 0, "size": 16}},
369
+ ],
370
+ )
371
+ n_workgroups = math.ceil(self.size / WORKGROUP_SIZE)
372
+ encoder = device.create_command_encoder()
373
+ pass_enc = encoder.begin_compute_pass()
374
+ pass_enc.set_pipeline(pipeline)
375
+ pass_enc.set_bind_group(0, bind_group)
376
+ pass_enc.dispatch_workgroups(n_workgroups)
377
+ pass_enc.end()
378
+ device.queue.submit([encoder.finish()])
379
+ uniform_buf.destroy()
380
+ return GPUArray(out_buf, self._shape)
381
+
382
+ # ----- argmax / argmin -----
383
+
384
+ def argmax(self) -> int:
385
+ """Return index of maximum value (as Python int)."""
386
+ data = self.to_numpy().flatten()
387
+ return int(np.argmax(data))
388
+
389
+ def argmin(self) -> int:
390
+ """Return index of minimum value (as Python int)."""
391
+ data = self.to_numpy().flatten()
392
+ return int(np.argmin(data))
393
+
394
+ # ----- copy -----
395
+
396
+ def copy(self) -> "GPUArray":
397
+ """Return a copy of this array."""
398
+ return self._run_unary_op(shaders.COPY_SHADER)
399
+
400
+ # ----- comparison operators -----
401
+
402
+ def __eq__(self, other):
403
+ if isinstance(other, (int, float)):
404
+ other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
405
+ return self._run_binary_op(other, shaders.EQ_SHADER)
406
+
407
+ def __gt__(self, other):
408
+ if isinstance(other, (int, float)):
409
+ other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
410
+ return self._run_binary_op(other, shaders.GT_SHADER)
411
+
412
+ def __lt__(self, other):
413
+ if isinstance(other, (int, float)):
414
+ other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
415
+ return self._run_binary_op(other, shaders.LT_SHADER)
416
+
417
+ def __ge__(self, other):
418
+ if isinstance(other, (int, float)):
419
+ other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
420
+ return self._run_binary_op(other, shaders.GE_SHADER)
421
+
422
+ def __le__(self, other):
423
+ if isinstance(other, (int, float)):
424
+ other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
425
+ return self._run_binary_op(other, shaders.LE_SHADER)
426
+
427
+ # ----- power -----
428
+
429
+ def __pow__(self, other):
430
+ if isinstance(other, (int, float)):
431
+ other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
432
+ return self._run_binary_op(other, shaders.POW_SHADER)
433
+
434
+ # ----- flatten -----
435
+
436
+ def flatten(self) -> "GPUArray":
437
+ """Return a 1D view of the array."""
438
+ return GPUArray(self._buffer, (self.size,))
439
+
440
+ # ----- len -----
441
+
442
+ def __len__(self) -> int:
443
+ return self._shape[0]
444
+
346
445
 
347
446
  def _from_numpy(data: np.ndarray) -> GPUArray:
348
447
  """Create a GPUArray from a numpy array."""
@@ -424,3 +424,177 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
424
424
  c[row * dims.N + col] = sum;
425
425
  }
426
426
  """
427
+
428
+ # ---------------------------------------------------------------------------
429
+ # Clamp (requires uniform buffer with min_val, max_val)
430
+ # ---------------------------------------------------------------------------
431
+
432
+ CLAMP_SHADER = """
433
+ struct Params {
434
+ min_val: f32,
435
+ max_val: f32,
436
+ }
437
+
438
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
439
+ @group(0) @binding(1) var<storage, read_write> result: array<f32>;
440
+ @group(0) @binding(2) var<uniform> params: Params;
441
+
442
+ @compute @workgroup_size(256)
443
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
444
+ let i = gid.x;
445
+ if (i < arrayLength(&a)) {
446
+ result[i] = clamp(a[i], params.min_val, params.max_val);
447
+ }
448
+ }
449
+ """
450
+
451
+ # ---------------------------------------------------------------------------
452
+ # Where / select: result = cond * a + (1-cond) * b
453
+ # ---------------------------------------------------------------------------
454
+
455
+ WHERE_SHADER = """
456
+ @group(0) @binding(0) var<storage, read> cond: array<f32>;
457
+ @group(0) @binding(1) var<storage, read> a: array<f32>;
458
+ @group(0) @binding(2) var<storage, read> b: array<f32>;
459
+ @group(0) @binding(3) var<storage, read_write> result: array<f32>;
460
+
461
+ @compute @workgroup_size(256)
462
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
463
+ let i = gid.x;
464
+ if (i < arrayLength(&cond)) {
465
+ result[i] = cond[i] * a[i] + (1.0 - cond[i]) * b[i];
466
+ }
467
+ }
468
+ """
469
+
470
+ # ---------------------------------------------------------------------------
471
+ # Comparison shaders (return 0.0 or 1.0)
472
+ # ---------------------------------------------------------------------------
473
+
474
+ EQ_SHADER = """
475
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
476
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
477
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
478
+
479
+ @compute @workgroup_size(256)
480
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
481
+ let i = gid.x;
482
+ if (i < arrayLength(&a)) {
483
+ if (a[i] == b[i]) {
484
+ result[i] = 1.0;
485
+ } else {
486
+ result[i] = 0.0;
487
+ }
488
+ }
489
+ }
490
+ """
491
+
492
+ GT_SHADER = """
493
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
494
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
495
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
496
+
497
+ @compute @workgroup_size(256)
498
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
499
+ let i = gid.x;
500
+ if (i < arrayLength(&a)) {
501
+ if (a[i] > b[i]) {
502
+ result[i] = 1.0;
503
+ } else {
504
+ result[i] = 0.0;
505
+ }
506
+ }
507
+ }
508
+ """
509
+
510
+ LT_SHADER = """
511
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
512
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
513
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
514
+
515
+ @compute @workgroup_size(256)
516
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
517
+ let i = gid.x;
518
+ if (i < arrayLength(&a)) {
519
+ if (a[i] < b[i]) {
520
+ result[i] = 1.0;
521
+ } else {
522
+ result[i] = 0.0;
523
+ }
524
+ }
525
+ }
526
+ """
527
+
528
+ GE_SHADER = """
529
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
530
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
531
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
532
+
533
+ @compute @workgroup_size(256)
534
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
535
+ let i = gid.x;
536
+ if (i < arrayLength(&a)) {
537
+ if (a[i] >= b[i]) {
538
+ result[i] = 1.0;
539
+ } else {
540
+ result[i] = 0.0;
541
+ }
542
+ }
543
+ }
544
+ """
545
+
546
+ LE_SHADER = """
547
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
548
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
549
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
550
+
551
+ @compute @workgroup_size(256)
552
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
553
+ let i = gid.x;
554
+ if (i < arrayLength(&a)) {
555
+ if (a[i] <= b[i]) {
556
+ result[i] = 1.0;
557
+ } else {
558
+ result[i] = 0.0;
559
+ }
560
+ }
561
+ }
562
+ """
563
+
564
+ # ---------------------------------------------------------------------------
565
+ # Fill (fill array with uniform scalar)
566
+ # ---------------------------------------------------------------------------
567
+
568
+ FILL_SHADER = """
569
+ struct Params {
570
+ scalar: f32,
571
+ }
572
+
573
+ @group(0) @binding(0) var<storage, read_write> result: array<f32>;
574
+ @group(0) @binding(1) var<uniform> params: Params;
575
+
576
+ @compute @workgroup_size(256)
577
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
578
+ let i = gid.x;
579
+ if (i < arrayLength(&result)) {
580
+ result[i] = params.scalar;
581
+ }
582
+ }
583
+ """
584
+
585
+ # ---------------------------------------------------------------------------
586
+ # Copy
587
+ # ---------------------------------------------------------------------------
588
+
589
+ COPY_SHADER = """
590
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
591
+ @group(0) @binding(1) var<storage, read_write> result: array<f32>;
592
+
593
+ @compute @workgroup_size(256)
594
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
595
+ let i = gid.x;
596
+ if (i < arrayLength(&a)) {
597
+ result[i] = a[i];
598
+ }
599
+ }
600
+ """
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: gpuarray
3
- Version: 0.1.0
3
+ Version: 0.2.0
4
4
  Summary: NumPy-like GPU array library that works on every GPU — NVIDIA, AMD, Intel, Apple Silicon. No CUDA required.
5
5
  Author: Arnav Kewalram
6
6
  License: MIT
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "gpuarray"
7
- version = "0.1.0"
7
+ version = "0.2.0"
8
8
  description = "NumPy-like GPU array library that works on every GPU — NVIDIA, AMD, Intel, Apple Silicon. No CUDA required."
9
9
  requires-python = ">=3.10,<3.14"
10
10
  license = {text = "MIT"}
@@ -225,3 +225,211 @@ class TestMisc:
225
225
  def test_getitem(self):
226
226
  a = gp.array([10, 20, 30])
227
227
  assert a[1] == 20.0
228
+
229
+
230
+ class TestClip:
231
+ def test_clip_basic(self):
232
+ a_np = np.array([-2, -1, 0, 1, 2, 3, 4, 5], dtype=np.float32)
233
+ a = gp.array(a_np)
234
+ result = a.clip(0.0, 3.0)
235
+ expected = np.clip(a_np, 0.0, 3.0)
236
+ npt.assert_allclose(result.to_numpy(), expected, rtol=RTOL)
237
+
238
+ def test_clip_api(self):
239
+ a_np = np.array([1, 5, 10, -3], dtype=np.float32)
240
+ a = gp.array(a_np)
241
+ result = gp.clip(a, 0.0, 7.0)
242
+ expected = np.clip(a_np, 0.0, 7.0)
243
+ npt.assert_allclose(result.to_numpy(), expected, rtol=RTOL)
244
+
245
+
246
+ class TestArgmaxArgmin:
247
+ def test_argmax(self):
248
+ a_np = np.array([1, 5, 3, 9, 2], dtype=np.float32)
249
+ a = gp.array(a_np)
250
+ assert a.argmax() == int(np.argmax(a_np))
251
+
252
+ def test_argmin(self):
253
+ a_np = np.array([3, 1, 4, 0, 5], dtype=np.float32)
254
+ a = gp.array(a_np)
255
+ assert a.argmin() == int(np.argmin(a_np))
256
+
257
+ def test_argmax_api(self):
258
+ a_np = np.array([10, 20, 5], dtype=np.float32)
259
+ a = gp.array(a_np)
260
+ assert gp.argmax(a) == 1
261
+
262
+ def test_argmin_api(self):
263
+ a_np = np.array([10, 20, 5], dtype=np.float32)
264
+ a = gp.array(a_np)
265
+ assert gp.argmin(a) == 2
266
+
267
+
268
+ class TestComparisons:
269
+ def test_eq(self):
270
+ a = gp.array([1.0, 2.0, 3.0])
271
+ b = gp.array([1.0, 5.0, 3.0])
272
+ result = (a == b).to_numpy()
273
+ expected = np.array([1.0, 0.0, 1.0], dtype=np.float32)
274
+ npt.assert_allclose(result, expected, rtol=RTOL)
275
+
276
+ def test_gt(self):
277
+ a = gp.array([1.0, 5.0, 3.0])
278
+ b = gp.array([2.0, 3.0, 3.0])
279
+ result = (a > b).to_numpy()
280
+ expected = np.array([0.0, 1.0, 0.0], dtype=np.float32)
281
+ npt.assert_allclose(result, expected, rtol=RTOL)
282
+
283
+ def test_lt(self):
284
+ a = gp.array([1.0, 5.0, 3.0])
285
+ b = gp.array([2.0, 3.0, 3.0])
286
+ result = (a < b).to_numpy()
287
+ expected = np.array([1.0, 0.0, 0.0], dtype=np.float32)
288
+ npt.assert_allclose(result, expected, rtol=RTOL)
289
+
290
+ def test_ge(self):
291
+ a = gp.array([1.0, 5.0, 3.0])
292
+ b = gp.array([2.0, 3.0, 3.0])
293
+ result = (a >= b).to_numpy()
294
+ expected = np.array([0.0, 1.0, 1.0], dtype=np.float32)
295
+ npt.assert_allclose(result, expected, rtol=RTOL)
296
+
297
+ def test_le(self):
298
+ a = gp.array([1.0, 5.0, 3.0])
299
+ b = gp.array([2.0, 3.0, 3.0])
300
+ result = (a <= b).to_numpy()
301
+ expected = np.array([1.0, 0.0, 1.0], dtype=np.float32)
302
+ npt.assert_allclose(result, expected, rtol=RTOL)
303
+
304
+ def test_eq_scalar(self):
305
+ a = gp.array([1.0, 2.0, 3.0])
306
+ result = (a == 2.0).to_numpy()
307
+ expected = np.array([0.0, 1.0, 0.0], dtype=np.float32)
308
+ npt.assert_allclose(result, expected, rtol=RTOL)
309
+
310
+ def test_gt_scalar(self):
311
+ a = gp.array([1.0, 2.0, 3.0])
312
+ result = (a > 1.5).to_numpy()
313
+ expected = np.array([0.0, 1.0, 1.0], dtype=np.float32)
314
+ npt.assert_allclose(result, expected, rtol=RTOL)
315
+
316
+
317
+ class TestWhere:
318
+ def test_where_basic(self):
319
+ cond = gp.array([1.0, 0.0, 1.0, 0.0])
320
+ x = gp.array([10.0, 20.0, 30.0, 40.0])
321
+ y = gp.array([50.0, 60.0, 70.0, 80.0])
322
+ result = gp.where(cond, x, y)
323
+ expected = np.array([10.0, 60.0, 30.0, 80.0], dtype=np.float32)
324
+ npt.assert_allclose(result.to_numpy(), expected, rtol=RTOL)
325
+
326
+ def test_where_from_comparison(self):
327
+ a = gp.array([1.0, 5.0, 3.0, 8.0])
328
+ cond = a > 3.0
329
+ x = gp.array([100.0, 100.0, 100.0, 100.0])
330
+ y = gp.array([0.0, 0.0, 0.0, 0.0])
331
+ result = gp.where(cond, x, y)
332
+ expected = np.array([0.0, 100.0, 0.0, 100.0], dtype=np.float32)
333
+ npt.assert_allclose(result.to_numpy(), expected, rtol=RTOL)
334
+
335
+
336
+ class TestConcatenate:
337
+ def test_concat_1d(self):
338
+ a = gp.array([1.0, 2.0, 3.0])
339
+ b = gp.array([4.0, 5.0, 6.0])
340
+ result = gp.concatenate([a, b])
341
+ expected = np.array([1, 2, 3, 4, 5, 6], dtype=np.float32)
342
+ npt.assert_allclose(result.to_numpy(), expected, rtol=RTOL)
343
+ assert result.shape == (6,)
344
+
345
+ def test_concat_2d(self):
346
+ a = gp.array([[1.0, 2.0], [3.0, 4.0]])
347
+ b = gp.array([[5.0, 6.0]])
348
+ result = gp.concatenate([a, b], axis=0)
349
+ expected = np.array([[1, 2], [3, 4], [5, 6]], dtype=np.float32)
350
+ npt.assert_allclose(result.to_numpy(), expected, rtol=RTOL)
351
+ assert result.shape == (3, 2)
352
+
353
+ def test_stack_1d(self):
354
+ a = gp.array([1.0, 2.0, 3.0])
355
+ b = gp.array([4.0, 5.0, 6.0])
356
+ result = gp.stack([a, b])
357
+ expected = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)
358
+ npt.assert_allclose(result.to_numpy(), expected, rtol=RTOL)
359
+ assert result.shape == (2, 3)
360
+
361
+
362
+ class TestPower:
363
+ def test_pow_array(self):
364
+ a_np = np.array([2.0, 3.0, 4.0], dtype=np.float32)
365
+ b_np = np.array([2.0, 2.0, 3.0], dtype=np.float32)
366
+ a = gp.array(a_np)
367
+ b = gp.array(b_np)
368
+ result = (a ** b).to_numpy()
369
+ expected = a_np ** b_np
370
+ npt.assert_allclose(result, expected, rtol=RTOL)
371
+
372
+ def test_pow_scalar(self):
373
+ a_np = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32)
374
+ a = gp.array(a_np)
375
+ result = (a ** 2).to_numpy()
376
+ expected = a_np ** 2
377
+ npt.assert_allclose(result, expected, rtol=RTOL)
378
+
379
+
380
+ class TestCopy:
381
+ def test_copy(self):
382
+ a = gp.array([1.0, 2.0, 3.0])
383
+ b = a.copy()
384
+ npt.assert_allclose(b.to_numpy(), a.to_numpy(), rtol=RTOL)
385
+ assert b._buffer is not a._buffer
386
+
387
+ def test_copy_shape(self):
388
+ a = gp.array([[1.0, 2.0], [3.0, 4.0]])
389
+ b = a.copy()
390
+ assert b.shape == a.shape
391
+ npt.assert_allclose(b.to_numpy(), a.to_numpy(), rtol=RTOL)
392
+
393
+
394
+ class TestRandom:
395
+ def test_rand_shape(self):
396
+ a = gp.random.rand(5)
397
+ assert a.shape == (5,)
398
+ data = a.to_numpy()
399
+ assert data.min() >= 0.0
400
+ assert data.max() < 1.0
401
+
402
+ def test_rand_2d(self):
403
+ a = gp.random.rand(3, 4)
404
+ assert a.shape == (3, 4)
405
+ assert a.size == 12
406
+
407
+ def test_randn_shape(self):
408
+ a = gp.random.randn(100)
409
+ assert a.shape == (100,)
410
+ data = a.to_numpy()
411
+ # randn should have mean near 0 and std near 1 for large samples
412
+ assert abs(data.mean()) < 0.5 # loose check
413
+
414
+ def test_randn_2d(self):
415
+ a = gp.random.randn(4, 5)
416
+ assert a.shape == (4, 5)
417
+ assert a.size == 20
418
+
419
+
420
+ class TestFlatten:
421
+ def test_flatten_2d(self):
422
+ a = gp.array([[1.0, 2.0], [3.0, 4.0]])
423
+ b = a.flatten()
424
+ assert b.shape == (4,)
425
+ npt.assert_allclose(b.to_numpy(), np.array([1, 2, 3, 4], dtype=np.float32), rtol=RTOL)
426
+
427
+
428
+ class TestLen:
429
+ def test_len_1d(self):
430
+ a = gp.array([1.0, 2.0, 3.0])
431
+ assert len(a) == 3
432
+
433
+ def test_len_2d(self):
434
+ a = gp.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
435
+ assert len(a) == 3
@@ -1,99 +0,0 @@
1
- """NumPy-compatible module-level functions."""
2
-
3
- from __future__ import annotations
4
-
5
- from typing import Sequence
6
-
7
- import numpy as np
8
-
9
- from .core import GPUArray, _from_numpy, matmul as _matmul
10
- from . import shaders
11
-
12
-
13
- # ----- creation functions -----
14
-
15
- def array(data) -> GPUArray:
16
- """Create a GPUArray from a list, nested list, or numpy array."""
17
- if isinstance(data, GPUArray):
18
- return data
19
- arr = np.asarray(data, dtype=np.float32)
20
- if arr.ndim == 0:
21
- arr = arr.reshape((1,))
22
- return _from_numpy(arr)
23
-
24
-
25
- def zeros(shape) -> GPUArray:
26
- if isinstance(shape, int):
27
- shape = (shape,)
28
- return _from_numpy(np.zeros(shape, dtype=np.float32))
29
-
30
-
31
- def ones(shape) -> GPUArray:
32
- if isinstance(shape, int):
33
- shape = (shape,)
34
- return _from_numpy(np.ones(shape, dtype=np.float32))
35
-
36
-
37
- def arange(start, stop=None, step=1) -> GPUArray:
38
- if stop is None:
39
- stop = start
40
- start = 0
41
- return _from_numpy(np.arange(start, stop, step, dtype=np.float32))
42
-
43
-
44
- def linspace(start, stop, num=50) -> GPUArray:
45
- return _from_numpy(np.linspace(start, stop, num, dtype=np.float32))
46
-
47
-
48
- # ----- operations -----
49
-
50
- def dot(a: GPUArray, b: GPUArray):
51
- return a.dot(b)
52
-
53
-
54
- def matmul(a: GPUArray, b: GPUArray) -> GPUArray:
55
- return _matmul(a, b)
56
-
57
-
58
- def sum(a: GPUArray) -> float:
59
- return a.sum()
60
-
61
-
62
- def mean(a: GPUArray) -> float:
63
- return a.mean()
64
-
65
-
66
- def exp(a: GPUArray) -> GPUArray:
67
- return a._run_unary_op(shaders.EXP_SHADER)
68
-
69
-
70
- def log(a: GPUArray) -> GPUArray:
71
- return a._run_unary_op(shaders.LOG_SHADER)
72
-
73
-
74
- def sqrt(a: GPUArray) -> GPUArray:
75
- return a._run_unary_op(shaders.SQRT_SHADER)
76
-
77
-
78
- def abs(a: GPUArray) -> GPUArray:
79
- return a._run_unary_op(shaders.ABS_SHADER)
80
-
81
-
82
- def max(a: GPUArray) -> float:
83
- return a.max()
84
-
85
-
86
- def min(a: GPUArray) -> float:
87
- return a.min()
88
-
89
-
90
- def relu(a: GPUArray) -> GPUArray:
91
- return a._run_unary_op(shaders.RELU_SHADER)
92
-
93
-
94
- def sigmoid(a: GPUArray) -> GPUArray:
95
- return a._run_unary_op(shaders.SIGMOID_SHADER)
96
-
97
-
98
- def tanh(a: GPUArray) -> GPUArray:
99
- return a._run_unary_op(shaders.TANH_SHADER)
File without changes
File without changes
File without changes