ML library for your hardware

• Eager-ish Execution — tensor operations fuse into kernels as you write them; when fusion is no longer possible, the kernel executes. For one off computations.
• Tape Mode — wrap loops in a Tape for lazy graph building, autograd and egraph-based fusion optimization. For repeated computations.
• Cross‑Platform Backends — codegen for C, CUDA, OpenCL and SPIR-V.
• Full 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 tape can be differentiated w.r.t. any other tensor in 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 automatcially. 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.

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.relu() * y
grads = tape.gradient(result, [x, y])

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

# from PyPI
pip install zyx-py

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

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

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.

flowchart LR
    A["Tensor E-Graph"] --> B["Fusion and Device Schedule Search"]
    B --> C["AOT kernels"]
    B --> D["Unified Kernel IR"]
    E --> F["IR autotuner with backend specific passes"]
    F --> G["Backend Code / Assembly"]

Outside a tape, tensor operations fuse eagerly into kernels as you call them. Inside a tape, a lazy graph is built and analyzed for fusion opportunities during realization or may pattern match parts of graph into AOT kernels. The fused operations are lowered to a unified intermediate representation, then compiled to native code for the target backend. Tape egraph compares fusion schemes and device allocations.

• 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 wgpu

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

• full tenstorrent coverage
• pattern matching for e-graph AOT kernels
• benchmarks + model bring-up

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

• 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

Star on GitHub | API Docs