gpuarray 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

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/api.py

@@ -97,3 +97,86 @@ def sigmoid(a: GPUArray) -> GPUArray:
 
 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()

gpuarray/core.py

@@ -343,6 +343,105 @@ class GPUArray:
             raise ValueError(f"Cannot reshape array of size {self.size} into shape {new_shape}")
         return GPUArray(self._buffer, new_shape)
 
+    # ----- clamp / clip -----
+
+    def clip(self, min_val: float, max_val: float) -> "GPUArray":
+        """Clamp values to [min_val, max_val]."""
+        device = get_device()
+        _, pipeline = _get_pipeline(shaders.CLAMP_SHADER)
+
+        out_buf = device.create_buffer(
+            size=self.size * 4,
+            usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST,
+        )
+        # Pack two f32 values, padded to 16 bytes
+        params_data = struct.pack("ff", min_val, max_val) + b"\x00" * 8
+        uniform_buf = device.create_buffer_with_data(
+            data=params_data,
+            usage=wgpu.BufferUsage.UNIFORM,
+        )
+        bind_group = device.create_bind_group(
+            layout=pipeline.get_bind_group_layout(0),
+            entries=[
+                {"binding": 0, "resource": {"buffer": self._buffer, "offset": 0, "size": self.size * 4}},
+                {"binding": 1, "resource": {"buffer": out_buf, "offset": 0, "size": self.size * 4}},
+                {"binding": 2, "resource": {"buffer": uniform_buf, "offset": 0, "size": 16}},
+            ],
+        )
+        n_workgroups = math.ceil(self.size / WORKGROUP_SIZE)
+        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()])
+        uniform_buf.destroy()
+        return GPUArray(out_buf, self._shape)
+
+    # ----- argmax / argmin -----
+
+    def argmax(self) -> int:
+        """Return index of maximum value (as Python int)."""
+        data = self.to_numpy().flatten()
+        return int(np.argmax(data))
+
+    def argmin(self) -> int:
+        """Return index of minimum value (as Python int)."""
+        data = self.to_numpy().flatten()
+        return int(np.argmin(data))
+
+    # ----- copy -----
+
+    def copy(self) -> "GPUArray":
+        """Return a copy of this array."""
+        return self._run_unary_op(shaders.COPY_SHADER)
+
+    # ----- comparison operators -----
+
+    def __eq__(self, other):
+        if isinstance(other, (int, float)):
+            other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
+        return self._run_binary_op(other, shaders.EQ_SHADER)
+
+    def __gt__(self, other):
+        if isinstance(other, (int, float)):
+            other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
+        return self._run_binary_op(other, shaders.GT_SHADER)
+
+    def __lt__(self, other):
+        if isinstance(other, (int, float)):
+            other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
+        return self._run_binary_op(other, shaders.LT_SHADER)
+
+    def __ge__(self, other):
+        if isinstance(other, (int, float)):
+            other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
+        return self._run_binary_op(other, shaders.GE_SHADER)
+
+    def __le__(self, other):
+        if isinstance(other, (int, float)):
+            other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
+        return self._run_binary_op(other, shaders.LE_SHADER)
+
+    # ----- power -----
+
+    def __pow__(self, other):
+        if isinstance(other, (int, float)):
+            other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
+        return self._run_binary_op(other, shaders.POW_SHADER)
+
+    # ----- flatten -----
+
+    def flatten(self) -> "GPUArray":
+        """Return a 1D view of the array."""
+        return GPUArray(self._buffer, (self.size,))
+
+    # ----- len -----
+
+    def __len__(self) -> int:
+        return self._shape[0]
+
 
 def _from_numpy(data: np.ndarray) -> GPUArray:
     """Create a GPUArray from a numpy array."""

gpuarray/shaders.py

@@ -424,3 +424,177 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
     c[row * dims.N + col] = sum;
 }
 """
+
+# ---------------------------------------------------------------------------
+# Clamp (requires uniform buffer with min_val, max_val)
+# ---------------------------------------------------------------------------
+
+CLAMP_SHADER = """
+struct Params {
+    min_val: f32,
+    max_val: f32,
+}
+
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read_write> result: array<f32>;
+@group(0) @binding(2) var<uniform> params: Params;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&a)) {
+        result[i] = clamp(a[i], params.min_val, params.max_val);
+    }
+}
+"""
+
+# ---------------------------------------------------------------------------
+# Where / select: result = cond * a + (1-cond) * b
+# ---------------------------------------------------------------------------
+
+WHERE_SHADER = """
+@group(0) @binding(0) var<storage, read> cond: array<f32>;
+@group(0) @binding(1) var<storage, read> a: array<f32>;
+@group(0) @binding(2) var<storage, read> b: array<f32>;
+@group(0) @binding(3) var<storage, read_write> result: array<f32>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&cond)) {
+        result[i] = cond[i] * a[i] + (1.0 - cond[i]) * b[i];
+    }
+}
+"""
+
+# ---------------------------------------------------------------------------
+# Comparison shaders (return 0.0 or 1.0)
+# ---------------------------------------------------------------------------
+
+EQ_SHADER = """
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+@group(0) @binding(2) var<storage, read_write> result: array<f32>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&a)) {
+        if (a[i] == b[i]) {
+            result[i] = 1.0;
+        } else {
+            result[i] = 0.0;
+        }
+    }
+}
+"""
+
+GT_SHADER = """
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+@group(0) @binding(2) var<storage, read_write> result: array<f32>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&a)) {
+        if (a[i] > b[i]) {
+            result[i] = 1.0;
+        } else {
+            result[i] = 0.0;
+        }
+    }
+}
+"""
+
+LT_SHADER = """
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+@group(0) @binding(2) var<storage, read_write> result: array<f32>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&a)) {
+        if (a[i] < b[i]) {
+            result[i] = 1.0;
+        } else {
+            result[i] = 0.0;
+        }
+    }
+}
+"""
+
+GE_SHADER = """
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+@group(0) @binding(2) var<storage, read_write> result: array<f32>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&a)) {
+        if (a[i] >= b[i]) {
+            result[i] = 1.0;
+        } else {
+            result[i] = 0.0;
+        }
+    }
+}
+"""
+
+LE_SHADER = """
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+@group(0) @binding(2) var<storage, read_write> result: array<f32>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&a)) {
+        if (a[i] <= b[i]) {
+            result[i] = 1.0;
+        } else {
+            result[i] = 0.0;
+        }
+    }
+}
+"""
+
+# ---------------------------------------------------------------------------
+# Fill (fill array with uniform scalar)
+# ---------------------------------------------------------------------------
+
+FILL_SHADER = """
+struct Params {
+    scalar: f32,
+}
+
+@group(0) @binding(0) var<storage, read_write> result: array<f32>;
+@group(0) @binding(1) var<uniform> params: Params;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&result)) {
+        result[i] = params.scalar;
+    }
+}
+"""
+
+# ---------------------------------------------------------------------------
+# Copy
+# ---------------------------------------------------------------------------
+
+COPY_SHADER = """
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read_write> result: array<f32>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&a)) {
+        result[i] = a[i];
+    }
+}
+"""

{gpuarray-0.1.0.dist-info → gpuarray-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpuarray
-Version: 0.1.0
+Version: 0.2.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.2.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+gpuarray/__init__.py,sha256=MB9k9oXPvA6wRCCZii72HtpRxjXPtkiT2RaJsiVY-wY,752
+gpuarray/api.py,sha256=PlYqtvY7jEV3xIneGdRoqXJjVOP3w2fnHyZgVP6fqis,4874
+gpuarray/core.py,sha256=bvBlQ9YSTRUcYawpnHpWnD-zJFdeMV16mv_rNFxYBGQ,18934
+gpuarray/device.py,sha256=q80VOrKhy1u5i4AMLUKhWVBwFqQjoJ4tNBm234Wpf-A,452
+gpuarray/shaders.py,sha256=LhS3vEZc08Mx_OPnsv06bT9gRPThuHIxBuWO0V2C_SI,16301
+gpuarray-0.2.0.dist-info/METADATA,sha256=8ZL33ZQ7eVBnm4GuvjLxNVCiVl39u2maLCpGXbuI24w,3132
+gpuarray-0.2.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
+gpuarray-0.2.0.dist-info/top_level.txt,sha256=iruYli4LmPBqhaZxTVvzT1bT0VeN6dbC0GE29lT1CbI,9
+gpuarray-0.2.0.dist-info/RECORD,,

gpuarray-0.1.0.dist-info/RECORD

@@ -1,9 +0,0 @@
-gpuarray/__init__.py,sha256=yDS9YEnJUsiPgAnftBGRQJqG2FwL6oRwvOI98p7WZu8,568
-gpuarray/api.py,sha256=tlIjtSCmCs-CYbZSAmW5XPxAxZmmSX3-wuGXI-6u1Zs,2090
-gpuarray/core.py,sha256=GIs-PE6Y7m0htoBv1B1hbGxImkmLFkP6z_Zo9Gm5Joc,15122
-gpuarray/device.py,sha256=q80VOrKhy1u5i4AMLUKhWVBwFqQjoJ4tNBm234Wpf-A,452
-gpuarray/shaders.py,sha256=7fuWEtPVVHXnJxFBeT5SoQ467FDCqM6fP14Wqd8r0HA,11411
-gpuarray-0.1.0.dist-info/METADATA,sha256=WZHK9RoABh502yZZ7s8U4h4vy4UGknpfOQ1duSeb-3I,3132
-gpuarray-0.1.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
-gpuarray-0.1.0.dist-info/top_level.txt,sha256=iruYli4LmPBqhaZxTVvzT1bT0VeN6dbC0GE29lT1CbI,9
-gpuarray-0.1.0.dist-info/RECORD,,