wefml 0.1.2

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
• Sequential and functional model building
• 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)

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())

# 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)

Sequential Model

import wefml

model = wefml.Model([
    wefml.Linear(128, use_bias=True),
    wefml.ReLU(),
    wefml.Linear(10, use_bias=True),
])

model.fit(labels, inputs, epochs=10, lr=0.01)
pred = model.predict(test_inputs)
model.summary()

Functional Model (Transformer Example)

The functional API uses .call() to chain layers with explicit control over data flow, which is needed for architectures like transformers where layers have multiple inputs.

import wefml

# wefml.use_gpu(True)

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

d_model = 64

embedding = wefml.Embedding(tok.english_vsize, d_model)
embedding_out = wefml.Embedding(tok.spanish_vsize, d_model)
mha = wefml.MHA(d_model, self_attention=True, num_heads=4, use_bias=True, use_mask=True)
cross_attn = wefml.MHA(d_model, self_attention=False, num_heads=4, use_bias=True, use_mask=True)
ffn = wefml.Linear(d_model, use_bias=True)
out = wefml.Linear(tok.spanish_vsize, use_bias=True)
norm = wefml.LayerNorm(2)
relu = wefml.ReLU()

In the C++ codebase, the transformer uses the functional API like this:

// Functional model: layers are called explicitly with .call()
ValLayer x = {nullptr, &enc_input};
x = embedding.call(x, training, gpu);

// self attention
ValLayer c = mha_input.call(x, x, x, training, gpu, {nullptr, &enc_mask});

// add and norm
temp = *c.val + *x.val;
x.val = &temp;
x = norm.call(x, training, gpu);

// feed forward
c = ffn.call(x, training, gpu);
c = relu.call(c, training, gpu);

// output layer
e = out.call(x_out, training, gpu);

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)

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

• Linear(units, use_bias, seed) -- dense layer, GPU variant auto-selected
• Conv2D(kh, kw, units, use_bias, seed) -- 2D convolution, GPU variant auto-selected
• MaxPool2D(kh, kw) -- max pooling, GPU variant auto-selected
• ReLU() -- rectified linear activation
• Sigmoid() -- sigmoid activation
• Flatten() -- flatten to 1D
• LayerNorm(axis, eps) -- layer normalization
• ReduceSum(axis, keepdims) -- sum reduction along axis
• Embedding(vocab_size, d_model, seed) -- token embedding lookup
• MHA(d_model, self_attn, num_heads, use_bias, use_mask, use_gpu) -- multi-head attention

Model

Sequential: wefml.Model(layers, use_gpu)

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

Functional: layers have a .call() method for explicit wiring (used in transformers and other non-sequential architectures).

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