Mojo AI Inference Basics
AI inference means running a trained neural network to produce predictions. Training teaches the network its weights. Inference applies those weights to new input data to get an output — a classification, a text completion, a generated image. Mojo is built specifically to make inference fast on real hardware without specialized frameworks.
What Inference Looks Like
Input data
│
▼
┌──────────────────────────────────────┐
│ Neural Network │
│ │
│ Layer 1 → Layer 2 → Layer 3 → ... │
│ (each layer: weights × input + bias)│
└──────────────────────────────────────┘
│
▼
Output (prediction)
Example:
Input: [photo of a cat]
Output: {"cat": 0.98, "dog": 0.01, "bird": 0.01}
The Core Operation: Linear Layer
Every neural network layer performs the same fundamental computation: multiply an input vector by a weight matrix and add a bias vector. In math notation: output = W × input + b. This is a matrix-vector product followed by an addition — both operations that Mojo handles extremely efficiently.
fn linear_layer(
input: UnsafePointer[Float32],
weights: UnsafePointer[Float32],
bias: UnsafePointer[Float32],
output: UnsafePointer[Float32],
in_size: Int,
out_size: Int
):
for o in range(out_size):
var sum: Float32 = bias[o]
for i in range(in_size):
sum += weights[o * in_size + i] * input[i]
output[o] = sum
Matrix × Vector computation:
W (2×3 weights): input (3): output (2):
┌──────────────────┐ ┌──┐ ┌────────────────────────┐
│ 0.5 0.3 0.2 │ × │x0│ = │ 0.5x0 + 0.3x1 + 0.2x2 │
│ 0.1 0.4 0.8 │ │x1│ │ 0.1x0 + 0.4x1 + 0.8x2 │
└──────────────────┘ │x2│ └────────────────────────┘
└──┘
Then add bias to each output element.
Vectorized Linear Layer
The naive implementation loops one element at a time. A vectorized version processes multiple elements per instruction using SIMD, achieving 4–16× higher throughput.
from memory import UnsafePointer
fn linear_layer_simd(
input: UnsafePointer[Float32],
weights: UnsafePointer[Float32],
bias: UnsafePointer[Float32],
output: UnsafePointer[Float32],
in_size: Int,
out_size: Int
):
alias width = 8
for o in range(out_size):
var acc = SIMD[DType.float32, width].splat(0.0)
var i = 0
while i + width <= in_size:
var w = SIMD[DType.float32, width].load(weights + o * in_size + i)
var x = SIMD[DType.float32, width].load(input + i)
acc += w * x
i += width
var sum = acc.reduce_add()
while i < in_size:
sum += weights[o * in_size + i] * input[i]
i += 1
output[o] = sum + bias[o]
Activation Functions
After the linear computation, a non-linear activation function transforms the output. Without activations, stacking layers does nothing — a stack of linear layers collapses into one linear layer. Activations give networks the ability to model complex, non-linear patterns.
ReLU (Rectified Linear Unit)
The simplest and most commonly used activation: pass positive values through unchanged, replace negatives with zero.
fn relu(x: Float32) -> Float32:
return x if x > 0.0 else 0.0
fn relu_inplace(data: UnsafePointer[Float32], n: Int):
for i in range(n):
if data[i] < 0.0:
data[i] = 0.0
ReLU diagram:
Input: [-2, -1, 0, 1, 2, 3]
Output: [ 0, 0, 0, 1, 2, 3]
│
Negative values → 0 └── Positive values unchanged
Softmax
Used in the final layer for classification. Converts raw scores (logits) into a probability distribution — all outputs sum to 1.0.
from math import exp
fn softmax(logits: UnsafePointer[Float32], output: UnsafePointer[Float32], n: Int):
# Find max for numerical stability
var max_val = logits[0]
for i in range(1, n):
if logits[i] > max_val:
max_val = logits[i]
# Compute exp(x - max) and sum
var total: Float32 = 0.0
for i in range(n):
output[i] = exp(logits[i] - max_val)
total += output[i]
# Normalize to sum to 1
for i in range(n):
output[i] /= total
Softmax example:
logits: [2.0, 1.0, 0.5]
exp(): [7.39, 2.72, 1.65] (sum = 11.76)
output: [0.63, 0.23, 0.14] (sum = 1.0 ✓)
Interpretation: 63% confidence in class 0,
23% in class 1, 14% in class 2.
Building a Two-Layer Network
from memory import UnsafePointer
from math import exp
fn two_layer_inference(
input: UnsafePointer[Float32],
w1: UnsafePointer[Float32], b1: UnsafePointer[Float32],
w2: UnsafePointer[Float32], b2: UnsafePointer[Float32],
hidden: UnsafePointer[Float32],
output: UnsafePointer[Float32],
in_size: Int,
hid_size: Int,
out_size: Int
):
# Layer 1: linear + ReLU
for h in range(hid_size):
var s: Float32 = b1[h]
for i in range(in_size):
s += w1[h * in_size + i] * input[i]
hidden[h] = s if s > 0.0 else 0.0 # ReLU
# Layer 2: linear (softmax applied separately)
for o in range(out_size):
var s: Float32 = b2[o]
for h in range(hid_size):
s += w2[o * hid_size + h] * hidden[h]
output[o] = s
Network diagram:
Input (in_size)
│
▼
[Linear W1] → [ReLU] → Hidden (hid_size)
│
▼
[Linear W2] → [Softmax] → Output (out_size)
Loading Weights from a File
fn load_weights(path: String, ptr: UnsafePointer[Float32], n: Int) raises:
with open(path, "rb") as f:
var raw = f.read()
# In practice: parse binary float bytes from raw
# This pattern works with NumPy .npy files or raw binary exports
print("Loaded", n, "weights from", path)
Inference Performance Tips
Technique | Benefit -------------------|------------------------------------------- SIMD inner loop | 8× throughput on float32 operations Tiled matmul | Keeps weights in L1/L2 cache Parallelise rows | Use all CPU cores for large layers Float16 weights | Half the memory bandwidth vs Float32 Fuse operations | Linear + ReLU in one pass = less memory I/O Pre-allocate bufs | No allocation overhead during inference
Key Takeaways
Inference runs a trained model's weights on new input to produce predictions. The fundamental computation is a linear layer: matrix-vector multiply plus bias. Vectorize the inner dot-product loop with SIMD for 4–16× speedup. Apply ReLU between hidden layers to add non-linearity. Apply Softmax on the final layer to convert scores to probabilities. Combine tiling, SIMD, and parallelization for production-grade inference throughput. Mojo's direct hardware access lets you match — and often exceed — the speed of hand-tuned C++ inference engines.