Standard Library: Tensor

The tensor module provides the core tensor types and operations for numerical computation.

Key Exports

  • Tensor<T, Shape> — The primary tensor type with static shape.
  • zeros, ones, full — Tensor constructors.
  • add, mul, matmul — Element-wise and matrix operations.
  • sum, mean, max — Reduction operations.
  • reshape, transpose, squeeze, expand_dims — Shape manipulation.
  • slice, gather, index — Indexing and selection.
  • conv2d — 2D convolution with stride and padding control.
  • fft, ifft, fft2d — Spectral operations (FFT/IFFT).

Example Usage

use tensor::{Tensor, zeros, ones};

fn main() {
    let a: Tensor<f32>[2, 3] = zeros();
    let b: Tensor<f32>[2, 3] = ones();
    let c = a + b;  // Broadcasting add
    print(c);
}

Reduction Operations

Reductions accept keyword arguments for axis selection and dimension retention:

// Sum along axes with keyword arguments
let total = tensor.sum(x, axis=[1, 2], keepdims=true);

// Mean reduction
let avg = tensor.mean(activations, axis=[0]);

// Positional arguments also supported
let s = tensor.sum(x, [1, 2]);

Shape Operations

// Reshape to new dimensions
let flat = tensor.reshape(x, [batch_size, 784]);

// Transpose with permutation
let t = tensor.transpose(x, perm=[0, 2, 1]);

// Add/remove dimensions
let expanded = tensor.expand_dims(x, axis=1);
let squeezed = tensor.squeeze(x, axis=[2]);

Indexing and Slicing

// Slice with start:end:stride
let window = tensor.slice(x, [0, 0, 0], [batch, seq_len, dim]);

// Gather along axis
let selected = tensor.gather(embeddings, indices, axis=0);

// Convolution
let out = tensor.conv2d(input, filter, stride=[1, 1], padding="same");

Spectral Operations

MIND provides first-class FFT primitives compiled to native code via cuFFT (CUDA), rocFFT (ROCm), vDSP (Metal), custom WGSL shaders (WebGPU), or WebNN graph ops. All transforms run at O(N log N) with zero runtime dispatch overhead.

fft(signal)

1D Fast Fourier Transform. Real input [N] returns complex [N/2+1, 2]. Complex input [N, 2] returns complex [N, 2].

ifft(spectrum)

Inverse 1D FFT with automatic 1/N normalization. Guarantees ifft(fft(x)) == x.

fft2d(signal)

2D FFT for image processing and spatial filtering. O(H * W * log(H * W)).

Backend Dispatch

BackendLibrary
CUDAcuFFT
ROCmrocFFT
MetalvDSP / Accelerate
WebGPUWGSL Cooley-Tukey
WebNNMLGraphBuilder ops (CPU/GPU/NPU)

FFT Example: Low-Pass Filter

use tensor::{fft, ifft, zeros, ones};

fn low_pass_filter(signal: Tensor, cutoff: i64) -> Tensor {
    let spectrum = fft(signal);
    let n = spectrum.shape[0];
    let mask = zeros([n, 2]);
    for i in 0..cutoff {
        mask[i] = ones([2]);
        mask[n - 1 - i] = ones([2]);
    }
    return ifft(spectrum * mask);
}