Metadata-Version: 2.4
Name: computeml
Version: 0.4.0
Summary: ML inference and training on any GPU through Vulkan compute shaders, from NumPy
Author: Jorge Cardona
License: MIT
Project-URL: Homepage, https://gitlab.com/cardona-oss/compute-ml
Project-URL: Issues, https://gitlab.com/cardona-oss/compute-ml/-/tree/main/docs/issues
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: numpy>=1.22
Provides-Extra: test
Requires-Dist: pytest; extra == "test"

# computeml

Run ML inference and training on any GPU from Python, through Vulkan compute shaders. You work with NumPy
arrays. The wheel ships the compiled library, so you only need NumPy and a Vulkan driver.

## Install

```bash
pip install computeml
```

```bash
uv add computeml
```

Requirements:

- Linux x86_64 (glibc 2.27 or newer) and Python 3.10 or newer.
- A Vulkan driver. NVIDIA, AMD and Intel GPU drivers include one. With no GPU, Mesa lavapipe
  runs on the CPU: `apt install mesa-vulkan-drivers`.

## Quickstart

```python
import numpy as np
import computeml as cml

ctx = cml.Context()
print(ctx.device_name)

a = ctx.asarray(np.random.rand(512, 256))
b = ctx.asarray(np.random.rand(256, 128))
c = cml.ops.relu(a @ b) * 2.0
print(c.shape, c.numpy()[0, :4])
```

`ctx.asarray` copies an array to the GPU. Ops queue work on the GPU. `.numpy()` waits for the
work and copies the result back.

## Coming from NumPy, PyTorch or CuPy

| You know | In computeml |
|---|---|
| `cp.asarray(x)`, `torch.as_tensor(x)` | `ctx.asarray(x)` |
| `np.zeros`, `ones`, `full`, `empty` | `ctx.zeros(shape)`, `ctx.ones`, `ctx.full`, `ctx.empty` |
| `a + b`, `a * 2`, `a @ b`, `-a` | the same |
| `.shape`, `.dtype`, `.ndim`, `.size` | the same |
| `t.cpu().numpy()`, `cp.asnumpy(t)` | `t.numpy()` |
| `torch.nn.functional.layer_norm` | `cml.ops.layer_norm` |
| `torch.nn.Linear(16, 32)` | `cml.nn.Linear(ctx, 16, 32)` |
| `loss.backward()`, `w.grad`, `torch.optim.AdamW` | the same, `cml.optim.AdamW` |

Differences you will meet:

- No broadcasting. Shapes must match element for element, or the op raises `ValueError`.
- float32 only for inputs and outputs. `asarray` converts ints and float64 to float32.
- `tensor + number` is not supported yet. `tensor * number` and `tensor / number` work.
- `np.asarray(tensor)` raises `TypeError`. Call `.numpy()`, so every GPU-to-CPU copy is visible.
- Every layer takes the `Context` first. There is no default device.

## Layers

```python
import numpy as np
import computeml as cml

ctx = cml.Context()
model = cml.nn.Sequential(
    cml.nn.Linear(ctx, 16, 32, activation="gelu"),
    cml.nn.Linear(ctx, 32, 8),
    cml.nn.LayerNorm(ctx, 8),
)
x = ctx.asarray(np.random.rand(4, 16))
print(model(x).numpy().shape)
```

`Linear` weights use the PyTorch layout `(out, in)`. Load trained weights with
`cml.nn.Linear.from_numpy(ctx, weight, bias)`.

## Training

```python
import numpy as np
import computeml as cml

ctx = cml.Context()
X = np.random.rand(128, 4)
y = (X.sum(axis=1) > 2).astype(np.float32)

model = cml.nn.Sequential(
    cml.nn.Linear(ctx, 4, 16, activation="relu"),
    cml.nn.Linear(ctx, 16, 2),
).requires_grad_()
opt = cml.optim.AdamW(model.parameters(), lr=0.05)

x, t = ctx.asarray(X), ctx.asarray(y)
for _ in range(100):
    opt.zero_grad()
    loss = cml.ops.cross_entropy(model(x), t)
    loss.backward()
    opt.step()
print(float(loss))
```

Gradients are computed on the GPU by the C library's autograd tape. See the
[training guide](https://compute-ml-0a9c30.gitlab.io/python/training) for which ops can
be trained.

## Fused ops

One kernel does the work of several ops, so the data is read from GPU memory once.

```python
import numpy as np
import computeml as cml

ctx = cml.Context()
x = ctx.asarray(np.random.rand(8, 64))
w = ctx.asarray(np.random.rand(128, 64))
b = ctx.asarray(np.random.rand(128))
y = cml.ops.linear(x, w, b, activation="gelu")  # gelu(x @ w.T + b), one kernel
print(y.shape)
```

`cml.ops` also has softmax, layer and RMS norm, flash attention (multi-head, causal), RoPE,
embeddings, permute, concat and reductions. See
[`ops.py`](https://gitlab.com/cardona-oss/compute-ml/-/blob/main/python/computeml/ops.py).

## Capture and replay

`ctx.capture` records a function's GPU work on the first call. Later calls resubmit the
recorded work with new input values, with no Python per op.

```python
import numpy as np
import computeml as cml

ctx = cml.Context()
w = ctx.asarray(np.random.rand(64, 64))

@ctx.capture
def step(x):
    return cml.ops.relu(x @ w)

for _ in range(3):
    y = step(np.random.rand(32, 64))
    print(y.numpy().sum())
print(step.record_count)  # 1: recorded once, replayed twice
```

The output is the same tensor on every call, so read it before the next call. Any other op
on the context in between makes the next call record again.

## Quantized weights

```python
import numpy as np
import computeml as cml

ctx = cml.Context()
x = ctx.asarray(np.random.rand(16, 256))
w = ctx.asarray(np.random.rand(256, 64)).quantize("Q8_0")
print((x @ w).numpy().shape)  # w is dequantized inside the matmul kernel
```

## Errors and logs

- `cml.NoDeviceError`: no Vulkan driver or GPU. The message says what to install.
- `ImportError` on `import computeml`: the Vulkan loader `libvulkan.so.1` is missing.
- `ValueError` / `TypeError`: a wrong shape or dtype, raised before any GPU work.
- Set `CML_LOG=1` to print the chosen GPU and kernel class to stderr.
- Set `CML_GPU_INDEX=1` to use the second GPU.

## Build from source

Set `COMPUTEML_LIB` to use your own `libcomputeml.so`. To build it and run the tests from a
checkout of the repo:

```bash
cmake -B build -S . && cmake --build build --target computeml_shared
cd python && uv run --with pytest python -m pytest -q
```

A BERT-tiny example lives in
[`examples/python/bert_tiny.py`](https://gitlab.com/cardona-oss/compute-ml/-/blob/main/examples/python/bert_tiny.py).
