gpuarray 0.1.0__tar.gz → 0.3.0__tar.gz

{gpuarray-0.1.0 → gpuarray-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpuarray
-Version: 0.1.0
+Version: 0.3.0
 Summary: NumPy-like GPU array library that works on every GPU — NVIDIA, AMD, Intel, Apple Silicon. No CUDA required.
 Author: Arnav Kewalram
 License: MIT

{gpuarray-0.1.0 → gpuarray-0.3.0}/gpuarray/__init__.py

@@ -20,6 +20,13 @@ from .api import (
     relu,
     sigmoid,
     tanh,
+    where,
+    clip,
+    argmax,
+    argmin,
+    concatenate,
+    stack,
+    random,
 )
 
 __version__ = "0.1.0"
@@ -44,4 +51,11 @@ __all__ = [
     "relu",
     "sigmoid",
     "tanh",
+    "where",
+    "clip",
+    "argmax",
+    "argmin",
+    "concatenate",
+    "stack",
+    "random",
 ]

gpuarray-0.3.0/gpuarray/api.py

@@ -0,0 +1,182 @@
+"""NumPy-compatible module-level functions."""
+
+from __future__ import annotations
+
+from typing import Sequence
+
+import numpy as np
+
+from .core import GPUArray, _from_numpy, matmul as _matmul
+from . import shaders
+
+
+# ----- creation functions -----
+
+def array(data) -> GPUArray:
+    """Create a GPUArray from a list, nested list, or numpy array."""
+    if isinstance(data, GPUArray):
+        return data
+    arr = np.asarray(data, dtype=np.float32)
+    if arr.ndim == 0:
+        arr = arr.reshape((1,))
+    return _from_numpy(arr)
+
+
+def zeros(shape) -> GPUArray:
+    if isinstance(shape, int):
+        shape = (shape,)
+    return _from_numpy(np.zeros(shape, dtype=np.float32))
+
+
+def ones(shape) -> GPUArray:
+    if isinstance(shape, int):
+        shape = (shape,)
+    return _from_numpy(np.ones(shape, dtype=np.float32))
+
+
+def arange(start, stop=None, step=1) -> GPUArray:
+    if stop is None:
+        stop = start
+        start = 0
+    return _from_numpy(np.arange(start, stop, step, dtype=np.float32))
+
+
+def linspace(start, stop, num=50) -> GPUArray:
+    return _from_numpy(np.linspace(start, stop, num, dtype=np.float32))
+
+
+# ----- operations -----
+
+def dot(a: GPUArray, b: GPUArray):
+    return a.dot(b)
+
+
+def matmul(a: GPUArray, b: GPUArray) -> GPUArray:
+    return _matmul(a, b)
+
+
+def sum(a: GPUArray) -> float:
+    return a.sum()
+
+
+def mean(a: GPUArray) -> float:
+    return a.mean()
+
+
+def exp(a: GPUArray) -> GPUArray:
+    return a._run_unary_op(shaders.EXP_SHADER)
+
+
+def log(a: GPUArray) -> GPUArray:
+    return a._run_unary_op(shaders.LOG_SHADER)
+
+
+def sqrt(a: GPUArray) -> GPUArray:
+    return a._run_unary_op(shaders.SQRT_SHADER)
+
+
+def abs(a: GPUArray) -> GPUArray:
+    return a._run_unary_op(shaders.ABS_SHADER)
+
+
+def max(a: GPUArray) -> float:
+    return a.max()
+
+
+def min(a: GPUArray) -> float:
+    return a.min()
+
+
+def relu(a: GPUArray) -> GPUArray:
+    return a._run_unary_op(shaders.RELU_SHADER)
+
+
+def sigmoid(a: GPUArray) -> GPUArray:
+    return a._run_unary_op(shaders.SIGMOID_SHADER)
+
+
+def tanh(a: GPUArray) -> GPUArray:
+    return a._run_unary_op(shaders.TANH_SHADER)
+
+
+# ----- new operations -----
+
+def where(condition: GPUArray, x: GPUArray, y: GPUArray) -> GPUArray:
+    """Conditional select: where condition is 1.0, pick x; else pick y."""
+    import wgpu
+    from .device import get_device
+    from .core import _get_pipeline
+
+    assert condition.size == x.size == y.size
+    device = get_device()
+    _, pipeline = _get_pipeline(shaders.WHERE_SHADER)
+
+    out_buf = device.create_buffer(
+        size=condition.size * 4,
+        usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST,
+    )
+    bind_group = device.create_bind_group(
+        layout=pipeline.get_bind_group_layout(0),
+        entries=[
+            {"binding": 0, "resource": {"buffer": condition._buffer, "offset": 0, "size": condition.size * 4}},
+            {"binding": 1, "resource": {"buffer": x._buffer, "offset": 0, "size": x.size * 4}},
+            {"binding": 2, "resource": {"buffer": y._buffer, "offset": 0, "size": y.size * 4}},
+            {"binding": 3, "resource": {"buffer": out_buf, "offset": 0, "size": condition.size * 4}},
+        ],
+    )
+    import math
+    n_workgroups = math.ceil(condition.size / 256)
+    encoder = device.create_command_encoder()
+    pass_enc = encoder.begin_compute_pass()
+    pass_enc.set_pipeline(pipeline)
+    pass_enc.set_bind_group(0, bind_group)
+    pass_enc.dispatch_workgroups(n_workgroups)
+    pass_enc.end()
+    device.queue.submit([encoder.finish()])
+    return GPUArray(out_buf, condition._shape)
+
+
+def clip(a: GPUArray, min_val: float, max_val: float) -> GPUArray:
+    """Clamp values to [min_val, max_val]."""
+    return a.clip(min_val, max_val)
+
+
+def argmax(a: GPUArray) -> int:
+    """Return index of maximum value."""
+    return a.argmax()
+
+
+def argmin(a: GPUArray) -> int:
+    """Return index of minimum value."""
+    return a.argmin()
+
+
+def concatenate(arrays, axis: int = 0) -> GPUArray:
+    """Join arrays along an axis (CPU-side, creates new GPU buffer)."""
+    np_arrays = [a.to_numpy() for a in arrays]
+    result = np.concatenate(np_arrays, axis=axis)
+    return _from_numpy(result)
+
+
+def stack(arrays, axis: int = 0) -> GPUArray:
+    """Stack arrays along a new axis (CPU-side, creates new GPU buffer)."""
+    np_arrays = [a.to_numpy() for a in arrays]
+    result = np.stack(np_arrays, axis=axis)
+    return _from_numpy(result)
+
+
+# ----- random submodule -----
+
+class _Random:
+    """Random array generation (CPU-side, uploaded to GPU)."""
+
+    def rand(self, *shape) -> GPUArray:
+        """Uniform random values in [0, 1)."""
+        return _from_numpy(np.random.rand(*shape).astype(np.float32))
+
+    def randn(self, *shape) -> GPUArray:
+        """Standard normal random values."""
+        return _from_numpy(np.random.randn(*shape).astype(np.float32))
+
+
+random = _Random()