wefml 0.1.0

A C++ deep learning library with Vulkan GPU acceleration, exposed to Python

wefml

C++ deep learning library with optional CUDA and Vulkan GPU acceleration, packaged for Python via pip.

Features

• Tensor class with full arithmetic and numpy interop
• Matrix ops: matmul, transpose, softmax, argmax, relu, sigmoid, reducesum, positional encoding
• Layers: Linear, Conv2D, MaxPool2D, Flatten, ReLU, Sigmoid, LayerNorm, ReduceSum, Embedding, Multi-Head Attention
• Model class: sequential model with .fit(), .predict(), .summary()
• Tokenizer: English/Spanish tab-separated tokenizer
• MNIST loader: IDX format image/label loader
• GPU toggle: one flag switches between CPU and GPU layers
• Dual GPU backend: CUDA (preferred) or Vulkan, auto-detected at build time

Installation

Prerequisites

• Python >= 3.8
• CMake >= 3.15
• A C++17 compiler (g++, clang++)
• pybind11 (pip install pybind11)
• (Optional) CUDA Toolkit for CUDA GPU support (preferred)
• (Optional) Vulkan SDK for Vulkan GPU support (fallback)

Install from source

cd learnn
pip install -e .

If neither CUDA nor Vulkan is found, the library builds in CPU-only mode automatically. GPU layers won't be available but everything else works.

Quick Start

import wefml
import numpy as np

# Check GPU availability
print(wefml.gpu_backend())         # 'cuda', 'vulkan', or 'none'
print(wefml.is_cuda_available())
print(wefml.is_vulkan_available())

# Enable GPU globally
# wefml.use_gpu(True)

# Tensors
a = wefml.Tensor.from_numpy(np.random.randn(2, 3).astype(np.float32))
b = wefml.Tensor.from_numpy(np.random.randn(3, 4).astype(np.float32))

c = wefml.ops.matmul(a, b)
print(c.numpy())

d = wefml.ops.transpose(a)
e = wefml.ops.softmax(c)
f = wefml.ops.relu(c)

# Build and train a model
model = wefml.Model([
    wefml.Linear(128, use_bias=True),
    wefml.ReLU(),
    wefml.Linear(10, use_bias=True),
])

model.summary()

# MNIST example
train_images = wefml.load_mnist_images("mnist/train-images-idx3-ubyte", max_items=1000)
train_labels = wefml.load_mnist_labels("mnist/train-labels-idx1-ubyte", max_items=1000)

test_images = wefml.load_mnist_images("mnist/t10k-images-idx3-ubyte", max_items=100)
test_labels = wefml.load_mnist_labels("mnist/t10k-labels-idx1-ubyte", max_items=100)

model = wefml.Model([
    wefml.Conv2D(3, 3, 16, use_bias=True),
    wefml.ReLU(),
    wefml.MaxPool2D(2, 2),
    wefml.Conv2D(3, 3, 8, use_bias=True),
    wefml.ReLU(),
    wefml.Flatten(),
    wefml.Linear(10, use_bias=True),
])

model.fit(
    train_labels, train_images,
    val_labels=test_labels, val_inputs=test_images,
    epochs=10, lr=0.01, batch_size=50
)

pred = model.predict(test_images)

GPU Mode

The build system auto-detects the GPU backend at compile time:

• CUDA is preferred when the CUDA toolkit is available (faster, NVIDIA GPUs)
• Vulkan is the fallback when CUDA is not found but Vulkan SDK is available (cross-platform)
• CPU-only when neither is available

Only one GPU backend is compiled since both implement the same class interfaces.

import wefml

print(wefml.gpu_backend())  # check which backend was compiled

wefml.use_gpu(True)

# Layer constructors automatically use GPU variants:
# Linear()    -> LinearGPU
# Conv2D()    -> Conv2DGPU
# MaxPool2D() -> MaxPool2DGPU

model = wefml.Model([
    wefml.Conv2D(3, 3, 128, use_bias=True),
    wefml.ReLU(),
    wefml.MaxPool2D(2, 2),
    wefml.Flatten(),
    wefml.Linear(10, use_bias=True),
], use_gpu=True)

Tokenizer Example

import wefml

tok = wefml.Tokenizer()
tok.process("english_spanish_tab.txt", early_stop=4000)

print(f"English vocab size: {tok.english_vsize}")
print(f"Spanish vocab size: {tok.spanish_vsize}")
print(f"Max sentence length: {tok.maxlen}")
tok.vocab_sample()

API Reference

GPU config

• wefml.use_gpu(enable: bool) -- toggle GPU acceleration globally
• wefml.is_gpu_available() -- True if any GPU support was compiled in
• wefml.is_cuda_available() -- True if CUDA was compiled in
• wefml.is_vulkan_available() -- True if Vulkan was compiled in
• wefml.gpu_backend() -- returns 'cuda', 'vulkan', or 'none'

Tensor

• Tensor.from_numpy(arr) -- create from numpy
• Tensor.zeros(shape) / Tensor.create(shape) -- create empty
• .numpy() -- convert to numpy
• .shape, .rank, .size
• +, -, *, / with other Tensors or scalars

Ops (wefml.ops)

• .matmul(a, b), .matmul_mt(a, b, threads)
• .transpose(a), .argmax(a), .softmax(a)
• .relu(a), .sigmoid(a), .reducesum(a, axis)
• .positional_encoding(length, depth)
• .l2(a, b), .binarycrossentropy(a, b)

Layers

Layer	CPU	GPU
Linear(units, use_bias, seed)	yes	yes
Conv2D(kh, kw, units, use_bias, seed)	yes	yes
MaxPool2D(kh, kw)	yes	yes
ReLU()	yes	yes
Sigmoid()	yes	yes
Flatten()	yes	yes
LayerNorm(axis, eps)	yes	yes
ReduceSum(axis, keepdims)	yes	yes
Embedding(vocab_size, d_model, seed)	yes	yes
MHA(d_model, self_attn, heads, ...)	yes	yes

Model (wefml.Model(layers, use_gpu))

• .add(layer)
• .fit(labels, inputs, val_labels, val_inputs, epochs, lr, batch_size, loss_fn)
• .predict(inputs)
• .summary()

Data

• wefml.Tokenizer() -- .process(path, early_stop)
• wefml.load_mnist_images(path, max_items)
• wefml.load_mnist_labels(path, max_items)

Building

The pip build uses scikit-build-core + pybind11 + CMake. When you run pip install ., it:

1. Compiles all C++ sources with pybind11 bindings
2. Detects CUDA first; if not found, detects Vulkan; links the available backend
3. Installs the wefml Python package with the native _learnn_core extension

License

MIT