Tensors & Kernels

View as Markdown

Neutron for Mojo provides type-safe tensors with compile-time dimension checking, backed by SIMD-accelerated kernels on the hot path. See the overview for the project's preview status.

Tensor Operations

from neutron.tensor import Tensor, Dim, matmul, softmax, rmsnorm

# Typed dimensions (compile-time shape safety)
alias Batch = Dim[0]
alias Seq = Dim[1]
alias Hidden = Dim[2]

# Core operations
let output = matmul(weights, input)      # Matrix multiply
let probs = softmax(logits, axis=-1)     # Softmax
let normed = rmsnorm(x, weight, 1e-5)    # RMS normalization
let activated = silu(x)                   # SiLU activation

SIMD Kernels

Hot-path operations use SIMD intrinsics for maximum throughput:

| Function | Description | |----------|-------------| | simd_dot(a, b) | Dot product | | simd_matvec(A, v) | Matrix-vector multiply | | simd_rmsnorm(x, w, eps) | RMS layer normalization | | simd_attention_scores(Q, K, scale) | Attention score computation | | simd_online_softmax_attention(Q, K, V) | Fused attention (FlashAttention-style) |

Additional Operations

layernorm(x, weight, bias)    # Layer normalization
gelu(x)                       # GELU activation
swiglu(x, w1, w2, w3)        # SwiGLU (used in LLaMA)