Zig ML (znn)

znn is a modular deep learning and LLM library written from scratch in pure Zig 0.17.0, with no runtime dependencies. The current version is recorded in build.zig.zon; release notes are in CHANGELOG.md.

• Tensors (src/tensor/): N-dimensional strided tensors, NumPy-style broadcasting, generic dtypes (f32, f64, bf16, i32, i64, usize, bool), zero-copy slices, unified operator option structs (ConvOptions, PoolOptions, RopeOptions), 1D/2D convolution and pooling kernels (src/tensor/conv_pool.zig), linear algebra decompositions (qr, symeig, svd in src/tensor/linalg.zig), and a compile-time shaped StaticTensor (src/tensor/static.zig).
• Autodiff (src/autodiff/): dynamic reverse-mode graph with operator-level topological sorting, arena-allocated activations and gradients, single-source forward kernels, initNoGrad inference mode, and module scopes used by the model graph export.
• Layers (src/nn/) & Models (src/models/): Linear (1D/2D/N-D), Conv1D, Conv2D, ConvTranspose1D, ConvTranspose2D, normalization layers (RMSNorm, LayerNorm, BatchNorm1d, BatchNorm2d, GroupNorm, Dropout), pooling layers (MaxPool1D, AvgPool1D, MaxPool2D, AvgPool2D, AdaptiveAvgPool1D, AdaptiveAvgPool2D), RNN / LSTM / StackedLSTM / GRU (RNNResult, LSTMResult, StackedLSTMResult, GRUResult), Transformer components (ScaledDotProductAttention, CausalSelfAttention with GQA / MQA, KVCache, MLALayer, SwiGLU, MoELayer, LoRALinear, GPT), Gemma 4 (Gemma4ForCausalLM, 4-bit Q4_0 quantized Gemma4Q4ForCausalLM, Gemma4Vocabulary), and a PyTorch-style module protocol built on one comptime field walker (walk, parameters, namedParameters, numParameters, setRequiresGrad, callForward, deinitModel, zeroGradModel, trainModel, evalModel, field-path naming via nameModules, forward scopes via enterModuleScope).
• Training (src/optim.zig, src/engine.zig, src/dataset.zig): SGD / Adam / AdamW, learning-rate schedulers, gradient clipping, Safetensors (F32/BF16/F16) models and binary optimizer checkpoints, SFT / DPO / GRPO losses, BPE tokenizer.
• Classical ML: OLS, Ridge, Lasso, ElasticNet (src/regression.zig), K-fold cross-validation (src/cross_validation.zig), t-SNE (src/manifold.zig).
• Hardware Acceleration & Multi-Backend (src/backend/): NVIDIA CUDA Driver API with zero-dependency PTX JIT compilation (GEMM, RMSNorm, Add, Mul, GELU, Softmax, RoPE), Apple Silicon MLX/Metal acceleration, macOS Accelerate (AMX) for GEMM, and a portable @Vector(8, f32) SIMD fallback.

Documentation

Topic	Location
Architecture: layers, tensor and autodiff internals, modules, optimizers, acceleration	doc/architecture.md
Model graph export (scopes, local graphs, JSON Schema)	doc/model-graph-visualization.md
Model graph visualization & edge routing specification	doc/visualization-model-edge-design.md
Interactive web visualizer (web/)	web/README.md
Documentation index	doc/README.md
Async execution & Apple Silicon Metal offload design (io: std.Io executor, Metal backend roadmap)	plan/ASYNC_OFFLOAD_DESIGN.md
Roadmap and task list	plan/TODO.md
Comparison with NumPy 2.x	plan/NUMPY_GAP_ANALYSIS.md
Contributor and agent rules	AGENTS.md

Quick Start

zig build test                          # all unit tests (includes check-layers)
zig build check-layers                  # verify one-way package dependencies under src/
zig build coverage                      # kcov coverage report (-- --fail-under=90, -- --open)
just bench                              # benchmark suite in ReleaseFast
just bench --suite gemm --filter conv   # suites: gemm, ops, activations, layers, models, optimizers, tokenizer

To run unit tests directly via the zig test CLI (e.g. for IDE integrations / ZLS) with Apple MLX and macOS system frameworks:

zig test src/root.zig \
  -I deps/mlx-c/include \
  -L deps/mlx-c/lib \
  -lmlxc -lmlx \
  -framework Metal -framework Foundation -framework QuartzCore -framework Accelerate \
  -lc

Datasets are downloaded on demand:

zig build download-dataset                     # Fashion MNIST (default)
zig build download-dataset -- mnist            # MNIST
zig build download-dataset -- tinyshakespeare  # also: wikitext2, tinystories, alpaca, all_llm

Examples

Command	Source	Description
zig build run -Doptimize=ReleaseFast	fashion_mnist.zig	3-layer MLP on Fashion MNIST
zig build run-cnn	cnn.zig	2D CNN on Fashion MNIST
zig build run-gan	gan.zig	GAN training with BceWithLogitsLoss
zig build run-lr	linear_regression.zig	OLS closed form vs. gradient descent
zig build run-logr	logistic_regression.zig	Logistic regression
zig build run-ridge	ridge_regression.zig	Ridge regression and multicollinearity
zig build run-reg	regularized_regression.zig	Ridge vs. Lasso vs. ElasticNet
zig build run-cv	cross_validation.zig	5-fold cross-validation grid search
zig build run-emb	transformer_embedding.zig	Token and position embeddings
zig build run-att	transformer_attention.zig	Causal multi-head self-attention
zig build run-block	transformer_block.zig	Transformer block forward step
zig build run-gpt	transformer_gpt.zig	GPT initialization and forward step
zig build run-shakespeare	train_shakespeare.zig	Mini-GPT on TinyShakespeare with text generation
zig build run-llm	llm_training.zig	BPE + SwiGLU + AdamW + SFT + LoRA + DPO pipeline
zig build run-report	export_model_report.zig	Model graph JSON export (sample_model_graph.json, minimal_model_graph.json)
zig build run-book-models	export_book_models.zig	Export Schema 2.0 JSON graphs for all 18 canonical book models (examples/models/*.json)
zig build run-static	comptime_static_tensor.zig	Compile-time shape-checked StaticTensor demonstration
zig build run-gemma4	gemma4.zig	Gemma 4 architecture and 4-bit (Q4_0) quantized inference demo (zig build quantize-gemma4 converts Safetensors weights); sampling defaults to the official generation config (temperature 1.0, top-k 64, top-p 0.95) and is tuned with --temperature, --top-k, --top-p, --penalty, --seed or --greedy
zig build run-bench	benchmark.zig	Benchmark CLI