zyx

ML library for your hardware

Why?

ML was enabled by new kinds of highly parallel, high performance hardware that did not exist before.

Zyx has 3 goals:

1. Be correct
2. Run everywhere (all hardware)
3. Run fast

And be nice to use.

ML won't get better without new hardware; existing libraries may not be the best fit for emerging hardware. The primary problem is the requirement to write custom kernels to get the required performance. Manufacturers have a tough time writing these high performance kernels; therefore they write kernels for only a few ops and don't support general linear algebra.

Zyx approaches these problems from two angles:

1. Supporting all ops

Zyx has linear SSA-ish IR with explicit control flow. This is the hardware unifying interface. Each piece of hardware has add instructions, can repeat instructions (loop), is highly parallel (work sizes), has multiple types of memory in a hierarchy (or at least global and registers) and can optionally have vectorization and tiling. This is the core of the instruction set. Zyx has a series of optimization passes (selected by autotuner) that apply various optimizations for different levels of these characteristics. The lowering layer from IR to backends is an almost 1:1 mapping. If your hardware can provide this translation, the whole stack (zyx ops, zyx-nn, zyx-optim) works on it.

2. Bringing top performance

As good as the automatic optimizations can be, writing manual kernels will usually be faster, which is why it is the dominant approach as of now. Zyx acknowledges this and its e-graph system pattern matching of any subgraph structure into a custom kernel written in a language of your choosing (or raw binary blobs, cublas, cblas, etc.), as well as writing custom kernels in zyx IR and taking advantage of optimization passes zyx provides. The e-graph measures their timings, compares them with auto-generated zyx kernels and picks the fastest path through this graph.

The other issue is running on edge platforms that don't have sufficient resources. Zyx takes only about 5 MB and uses machine-available drivers to run, such as a provided C compiler or CUDA runtime; for example, if you don't install CUDA, any GPU driver with Vulkan support is sufficient.

Features

• Eager mode, Lazy JIT Execution — tensor operations fuse into kernels as you write them; when fusion is no longer possible, the kernel executes. For one-off computations.
• Tape (e-graph) — wrap loop bodies in a Tape for lazy graph building, autograd and e-graph-based fusion optimization. Computation happens when realize is called. For repeated computations.
• Cross‑Platform Backends — codegen for C, CUDA, PTX, OpenCL and SPIR-V (Vulkan/WGPU).
• Linear‑Algebra Coverage — mirrors the PyTorch ops API (matmul, convolutions, pooling, reductions, indexing, etc.) by stacking ops. Stack more ops yourself to get more op coverage; zyx auto fuses and optimizes it.
• Immutable Tensors — tensors cannot be modified in place, preventing back‑prop errors common in PyTorch (RuntimeError: a tensor was modified in place).
• Explicit Tape — you control what is recorded via Tape; no need for torch.no_grad() or requires_grad semantics.
• Everything is diff — every tensor in a tape can be differentiated w.r.t. any other tensor in a tape.
• Lazy Device Loading — tensors load from their current memory pool (disk, another device) into the compute device only when needed.
• Parallel Pipelining — kernels allocate across heterogeneous devices (GPU, CPU, WebGPU) in a pipelined fashion via the scheduler automatically. e-graph tries all options, picks the fastest measured path.
• Small Footprint — compiled library is only a few MB with two dependencies (libloading, nanoserde) and std. This means for all models, a few-MB binary runs (and trains) them on all backends. Training and deployment can freely use the same API.

Crates

Installation

# from crates.io
cargo add zyx zyx-nn zyx-optim

# from PyPI, contains all backends, nn and optim
pip install zyx-py

Configuration & Debugging

• Configuration - Hardware device selection, autotune settings
• Environment Variables - Debug flags

🐍 Python Bindings

import zyx

x = zyx.Tensor.randn(2, 3)
y = zyx.Tensor.uniform_(2, 3, from_=-1.0, to_=1.0)
z = x.relu() + y.tanh()
print(z.shape())

# Autograd with tape
tape = zyx.Tape([x, y])
result = x.gelu() * y
grads = tape.gradient(result, [x, y])

Neural Nets

A training loop with a two-layer network, using Tape for autograd and optimizations:

use zyx::{Tensor, DType, Tape};
use zyx_nn::{Linear, Module};
use zyx_optim::SGD;

#[derive(Module)]
struct SimpleNet {
    linear1: Linear,
    linear2: Linear,
}

impl SimpleNet {
    fn new(dtype: DType) -> Result<Self, zyx::ZyxError> {
        Ok(Self {
            linear1: Linear::new(784, 128, true, dtype)?,
            linear2: Linear::new(128, 10, true, dtype)?,
        })
    }
    
    fn forward(&self, x: &Tensor) -> Tensor {
        let x = self.linear1.forward(x).unwrap().relu();
        self.linear2.forward(&x).unwrap()
    }
}

fn main() -> Result<(), zyx::ZyxError> {
    let mut model = SimpleNet::new(DType::F32)?;
    let mut optim = SGD::default();
    let x = Tensor::randn([64, 784], DType::F32)?;
    let target = Tensor::randn([64, 10], DType::F32)?;
    
    for epoch in 0..100 {
        let tape = Tape::new(&model)?;
        let output = model.forward(&x);
        let loss = output.mse_loss(&target)?;
        let grads = tape.gradient(&loss, &model);
        optim.update(&mut model, grads);
        tape.realize(&model)?;
    }
    
    Ok(())
}

For more complex examples:

• Examples - MNIST, RNN and others

Custom Kernels

Hand-optimize kernels for peak performance using hardware-specific features (e.g. tensor cores) using zyx IR:

use zyx::kernel::{Kernel, Scope, MemLayout, DeviceId};
use zyx::{DType, Tensor};

fn main() -> Result<(), zyx::ZyxError> {
    let mut kernel = Kernel::new(DeviceId::AUTO);
    let n = 4;
    let inp = kernel.define(DType::F32, Scope::Global, true, n);
    let gidx = kernel.gidx(0, n);
    let loaded = kernel.load(inp, gidx, MemLayout::Scalar);
    let doubled = kernel.add(loaded, loaded);
    let out = kernel.define(DType::F32, Scope::Global, false, n);
    kernel.store(out, doubled, gidx, MemLayout::Scalar);

    let compiled = kernel.compile()?;
    let x = Tensor::from([1.0f32, 2.0, 3.0, 4.0]);
    let result = compiled.forward(&[&x], [n]);
    let data: Vec<f32> = result.try_into().unwrap();
    assert_eq!(data, vec![2.0, 4.0, 6.0, 8.0]);
    Ok(())
}

See the WMMA matmul example for a tensor-core matmul example.

Architecture

graph TD
    A["Tensor ops"] --> B["Eager mode"]
    A --> C["Tape (e-graph)"]
    C --> D["Autograd"]
    D --> C
    C --> E["AOT kernels, fusion and device schedule search"]
    B --> F["Unified Kernel IR"]
    E --> F
    F --> G["IR autotuner with backend specific passes"]
    G --> H["Backend Code / Assembly"]

Outside the tape, tensor operations fuse eagerly into kernels as you call them using a unified kernel IR. Inside a tape, a lazy graph is built and analyzed for fusion opportunities during realization or may have parts pattern-matched into AOT kernels. Different device allocations are also compared. The fused operations are lowered to a unified kernel IR. Kernel IR is then autotuned and compiled to native code for the target backend.

How zyx compares

Zyx is a library, not a workflow: it doesn't prescribe training loops or data pipelines. The table below compares its design choices feature by feature.

Backends

• C - C codegen (clang/gcc)
• CUDA
• OpenCL
• Vulkan - SPIR-V codegen
• WGPU - SPIR-V codegen, feature: wgpu
• tenstorrent - Preliminary support, does not pass full test suite yet, feature tenstorrent

If you'd like to add a new backend to zyx, that would be awesome! Please read ADDING_BACKENDS.md

Benchmarks

Benchmarks are here in BENCHMARKS.md More will be added later. Feel free to try writing models in zyx and creating a PR with your your measured results.

Roadmap

• full tenstorrent coverage
• automatic device sharding search
• more backends
• more optimization passes
• more AOT kernels
• more benchmarks
• more model examples

Status & License

• Status: Stable API with active performance optimization
• License: LGPL-3.0-only (all crates)
• Rust Version: stable rust >= 1.88.0
• Platforms: Linux (primary), macOS, Windows (planned)

For Devs

• Architecture Book - How zyx works under the hood
• Contributing - How to contribute
• Adding backends - How to add new backends, information for hardware vendors
• Style - Zyx code style
• API Reference - Complete API documentation
• Issues - Bug reports and feature requests

---