AxonML v0.4.1

Performance Overhaul (v0.4.1)

Comprehensive performance optimization across all compute crates. 1,988 tests passing.

Backward Pass — BLAS Acceleration

• Conv2d/Conv1d/GroupedConv2d backward: im2col + GEMM replaces 7-deep nested loops (12x speedup, 107ms → 8.7ms on 64×64×32→64 conv)
• BatchNorm2d/1d backward: x_hat cache eliminates full recomputation in second pass
• AvgPool backward: Analytical count replaces double kernel loop
• Softmax backward: Precomputed outer_dim_strides eliminates O(N²) stride recomputation
• ConvTranspose2d backward: Cache-friendly loop reordering + zero-input skip

RNN/LSTM/GRU — Batched Matrix Multiply

• Pre-compute all input-to-hidden projections in one GEMM (1 large matmul vs seq_len small ones)
• GRU: Eliminated ones tensor allocation per timestep (n + z*(h-n) rewrite)
• LSTM: Only hidden-to-hidden matmul runs per timestep

Optimizers — Fused Update Loops

• Adam/AdamW: 3 separate loops → single pass, eliminated intermediate Vec allocations
• SGD: In-place momentum + parameter updates, eliminated redundant clones
• RMSprop: Fused denominator computation with param update
• LAMB: Fused update direction + both L2 norm computations into single pass

Tensor Operations

• Tensor::cat: copy_from_slice memcpy replaces element-by-element loop
• Tensor::var_dim: E[x²] - E[x]² formula (3 passes vs 4, saves one full-size allocation)

LLM Inference

• RMSNorm: Weight vector hoist outside batch loop
• RoPE: narrow() slice of cached cos/sin instead of full copy
• Causal mask: Half-loop (j in (i+1)..seq_len) instead of full conditional loop

Vision

• BlazeFace: Dynamic feature map sizes (fixes shape mismatch on non-128×128 inputs)
• NMS: Inner loop starts at pos+1 instead of scanning all
• Positional encoding: Precomputed frequency table

Bug Fixes

• Fixed JWT test failures (jsonwebtoken v10 rust_crypto feature)
• Fixed server integration tests (graceful skip on uninitialized DB)
• Fixed BlazeFace coco_bench_face_detectors shape mismatch

Full Changelog

See CHANGELOG.md for complete details.

Full Changelog: v0.3.0...v0.4.1