GPU Neural Networks via Vulkan

Learn Grilly

Build and train neural networks on any GPU — AMD, NVIDIA, or Intel — with a PyTorch-like API powered by Vulkan compute shaders.

pip install grilly
🧩

Tensors and Operations

Create arrays, index, reshape, broadcast, and run matrix math with numpy and VulkanTensor.

GPU Acceleration

Initialize Vulkan, enable GPU mode on modules, and use VulkanTensor for zero-copy compute.

🧠

Neural Networks

Define modules, choose loss functions and optimizers, and train an end-to-end model.

📦

Datasets & DataLoaders

Wrap data, split train/val, apply transforms, and iterate batches efficiently.

🖼

Convolutional Neural Networks

Build image classifiers with Conv2d, pooling, batch norm, and GPU-accelerated training.

🔌

Spiking Neural Networks

Train biologically-inspired LIF/IF neurons with surrogate gradients and temporal coding.

Unique to Grilly

Advanced Features

Flash attention, RoPE, LoRA fine-tuning, autograd Variables, and the functional API.

Quick Example

python
import grilly.nn as nn
import grilly.optim as optim
import numpy as np

model = nn.Sequential(
    nn.Linear(784, 256), nn.ReLU(),
    nn.Linear(256, 10),
)
x = np.random.randn(32, 784).astype(np.float32)
output = model(x)
print(output.shape)  # (32, 10)
Output
(32, 10)

Architecture

backend/ — Vulkan GPU dispatch, shader loading, pipeline caching
nn/ — PyTorch-like Module classes: Linear, Conv2d, LSTM, LIFNode, Transformer
functional/ — Stateless ops: F.relu, F.linear, F.flash_attention2, F.lif_step
optim/ — Optimizers: Adam, AdamW, SGD, AutoHypergradientAdamW
utils/ — Dataset, DataLoader, checkpointing, VulkanTensor, HuggingFace bridge