numr

Backends

• Kernels: SIMD on CPU (AVX-512, AVX2, NEON), native CUDA kernels, and WGSL compute shaders. No cuBLAS, MKL, or other vendor library.
• ROCm: not implemented.
• Dtype coverage: the table lists what a backend holds. Check the operation before you rely on a dtype.
• CUDA binaries: a build targets the GPUs on the build host. Set NUMR_CUDA_ARCH=all for a binary that runs on another machine.
• Details: the dtype table, SIMD, graph capture, and every NUMR_CUDA_ARCH value are in docs/backends.md. The CUDA build steps are in CUDA Build.

Installation

Needs Rust 1.95 or newer.

cargo add numr                   # CPU
cargo add numr --features cuda   # NVIDIA GPUs
cargo add numr --features wgpu   # WebGPU

Other feature flags:

What numr Is

numr provides:

• Tensors: N-dimensional arrays
• Linear algebra: decompositions and solvers
• FFT, statistics, and random distributions
• Autograd

numr is not a deep learning framework, a high-level ML API, or domain-specific.

Quick Start

CPU

use numr::prelude::*;
use numr::runtime::cpu::{CpuDevice, CpuRuntime};

fn main() -> Result<()> {
    let device = CpuDevice::new();

    let a = Tensor::<CpuRuntime>::from_slice(&[1.0_f32, 2.0, 3.0, 4.0], &[2, 2], &device)?;
    let b = Tensor::<CpuRuntime>::from_slice(&[5.0_f32, 6.0, 7.0, 8.0], &[2, 2], &device)?;

    // Arithmetic with broadcasting
    let c = a.add(&b)?;
    let d = a.mul(&b)?;

    // Matrix multiplication
    let e = a.matmul(&b)?;

    // Reductions: dims, keepdim
    let sum = c.sum(&[0, 1], false)?;
    let max = c.max(&[0, 1], false)?;

    // Element-wise functions
    let exp = a.exp()?;
    let sqrt = a.sqrt()?;

    // Reshape and transpose
    let flat = c.reshape(&[4])?;
    let transposed = c.transpose(0, 1)?;

    println!("{:?}", c.to_vec::<f32>());
    Ok(())
}

Backend-Generic Code

Write functions over R: Runtime. The same function runs on any backend.

use numr::prelude::*;

fn matrix_operations<R>(a: &Tensor<R>, b: &Tensor<R>, client: &R::Client) -> Result<Tensor<R>>
where
    R: Runtime,
    R::Client: BinaryOps<R> + MatmulOps<R> + ReduceOps<R>,
{
    let c = client.add(a, b)?;
    let d = client.matmul(&c, a)?;
    client.sum(&d, &[0, 1], false)
}

GPU (CUDA)

use numr::ops::RandomOps;
use numr::prelude::*;
use numr::runtime::cuda::{CudaClient, CudaRuntime};

fn main() -> std::result::Result<(), Box<dyn std::error::Error>> {
    let device = CudaRuntime::default_device();
    let client = CudaClient::new(device.clone())?;

    let a = client.randn(&[1024, 1024], DType::F32)?;
    let b = client.randn(&[1024, 1024], DType::F32)?;

    // Runs on the GPU through native CUDA kernels
    let c = client.matmul(&a, &b)?;

    // Copy the result to the host
    let host = c.try_to_vec::<f32>()?;
    Ok(())
}

The module docs cover the rest of the API:

• Linear algebra, FFT, and statistics: Algorithms and Operations
• Gradients: Autograd

Operations

numr provides tensor operations as traits in numr::ops. The same traits cover CPU, CUDA, and WebGPU. The full list is in docs/operations.md:

• Core arithmetic: unary, binary, scalar, type conversion, utility
• Shape and data movement: cat, stack, gather, scatter, sort, topk, and more
• Reductions, cumulative operations, comparisons, and logic
• Activations, normalization, convolution, einsum, and matmul with fused bias and activation
• Linear algebra and complex operations
• Statistics, random distributions, and quasi-random sequences
• Distance metrics on tensors and on plain slices

The algorithm modules are in docs/algorithms.md:

• Linear algebra decompositions, matrix functions, and solvers
• FFT and real FFT
• Special functions and polynomials
• Iterative solvers and preconditioners
• Sparse tensors (CSR, CSC, COO) and sparse linear algebra

For NumPy Users

Most numr operations are methods on the backend client, such as client.sum(&a, &[0], false). Import numr::prelude::*. The traits named below are in the prelude unless a row says otherwise.

NumPy has no equivalent for autograd, CUDA and WebGPU execution, or the fused ML operations. Those are in Operations.

Documentation

Building from Source

cargo build --release                    # CPU (default)
cargo build --release --features cuda    # CUDA
cargo build --release --features wgpu    # WebGPU
cargo test --release                     # Tests
cargo bench                              # Benchmarks

Contributing

Contributions are welcome. See CONTRIBUTING.md for conventions and checks.

Questions and design discussion happen on Discord.

License

Licensed under Apache-2.0. See LICENSE.