burn-flex vs burn-ndarray: Comprehensive Comparison

This document compares burn-flex (proposed replacement) against burn-ndarray (current CPU backend) to demonstrate full coverage and the architectural differences between the two.

Executive Summary

burn-flex is a from-scratch CPU backend built to replace burn-ndarray. The ndarray crate has been slow to evolve: it lacks f16/bf16 support, is limited to 6 dimensions, uses unsigned-only strides (preventing zero-copy flip), and simulates quantization rather than executing natively. burn-flex addresses all of these while passing the full burn-backend-tests suite, all ONNX model checks, and real model inference (ALBERT, MiniLM).

Performance improvements fall into two categories:

• Compute gains (1.1-9.7x): Better algorithms and libraries (gemm over matrixmultiply, Arc COW for buffer reuse, SIMD reductions).
• Structural improvements (up to 166,000x): Operations that burn-ndarray eagerly materializes (unfold, expand, slice, dequantize) are represented as zero-copy views or direct lookups in burn-flex, avoiding the work entirely.

burn-flex uses significantly less memory, supports f16/bf16 natively, runs on no_std/WASM/embedded, and has no dimension limit.

1. Architecture

Tensor Representation

FlexTensor (44 bytes without shape vec):

struct FlexTensor {
    data: Arc<Bytes>,    // 8 bytes (pointer)
    layout: Layout,      // shape + strides + offset
    dtype: DType,        // 1 byte enum
}

NdArrayTensor (enum with 11 typed variants):

enum NdArrayTensor {
    F64(NdArrayStorage<f64>),
    F32(NdArrayStorage<f32>),
    // ... 9 more variants
}

Key insight: Flex uses one struct for all dtypes with runtime dispatch. NdArray uses a typed enum with macro-based dispatch. Flex's approach is simpler (no macros, no generics plumbing) and enables operations to handle all dtypes uniformly.

Backend Type

Flex eliminates generic parameters entirely. Users write Flex instead of NdArray<f32, i64, i8>. Dtype selection happens at runtime via DType.

2. Feature Coverage

Float Dtypes

burn-flex's f16 support is native for all operations. For matmul and convolution, the gemm crate has native f16 kernels (since v0.15). bf16 converts to f32 for compute-heavy ops (matmul, conv) because gemm lacks native bf16 support.

Integer Dtypes

Both backends support the same integer dtypes.

Bool

Quantization

The fundamental difference is scale storage. Flex stores scales separately so dequantization is a simple scale * x_q multiply. NdArray stores everything in QuantizedBytes which must be parsed on every access, making it the bottleneck for all quantized operations.

3. Operation Coverage

Tensor Operations (FloatTensorOps)

All operations listed below are implemented by both backends unless marked otherwise.

Module Operations (ModuleOps)

Int and Bool Operations

Both backends implement all IntTensorOps and BoolTensorOps. The operations mirror float ops where applicable (arithmetic, comparison, reduction, gather/scatter, slice, etc.) plus type-specific operations (int_random uniform, bool_not, bool_and, bool_or, bool_xor).

Quantized Operations (QTensorOps)

Both backends implement all QTensorOps. The ops follow a dequantize-op-requantize pattern for most operations. Flex optimizes by:

• Storing scales separately for O(1) dequantization access
• Zero-copy layout ops on quantized tensors (permute, flip, expand, slice, select)
• Skipping dequantization for ordering ops (argmax, argmin, gather with tensor-level quant)

Activation Operations (ActivationOps)

Both backends implement all ActivationOps via the default trait implementations (relu, gelu, etc.).

Transaction Operations

Both backends implement TransactionOps for batched tensor operations.

4. Dimension Limits

burn-ndarray's dimension limit comes from its reshape! macro which matches on dimensions 1-6:

match $D {
    1 => reshape!(ty $ty, n 1, ...),
    // ...
    6 => reshape!(ty $ty, n 6, ...),
    _ => panic!("NdArray supports arrays up to 6 dimensions"),
}

burn-flex uses IxDyn-equivalent dynamic shapes with no upper bound.

5. Zero-Copy Operations

Flex's signed strides (isize) enable zero-copy flip, which is impossible with ndarray's unsigned strides. The unfold operation is especially dramatic: Flex returns a strided view in ~50ns regardless of size, while NdArray copies all window data (milliseconds for large tensors).

6. Memory Strategy

In-Place Mutation

Zero-Copy Loading

Both backends support zero-copy loading from external sources (burnpack files, mmap'd data):

7. SIMD

Both use macerator for portable SIMD. NdArray additionally has SIMD-optimized conv and pool kernels. Flex relies on the gemm crate's built-in SIMD for matmul/conv performance.

8. Matrix Multiplication

burn-ndarray offers optional BLAS acceleration (Accelerate on macOS, OpenBLAS, Netlib) through feature flags. burn-flex uses only the gemm crate, which is pure Rust but highly optimized with its own SIMD kernels. The gemm crate consistently outperforms matrixmultiply by 1.3-3.4x on Apple M3 Max.

9. Convolutions

Flex's unified 3D approach means one implementation covers all dimensionalities. The tradeoff is that 1D/2D convolutions expand dimensions (negligible overhead since gemm dominates).

NdArray has dedicated SIMD conv/pool kernels via macerator, which can be faster for specific patterns. Flex relies on the gemm crate's SIMD for all compute-heavy paths.

10. Parallelism

11. Platform Support

burn-flex has been explicitly tested on embedded targets with Burn's burn-no-std-tests integration suite (MNIST model inference).

12. Dependencies

burn-flex

Total: 7 required + 3 optional

burn-ndarray

Total: 12 required + 9 optional + 2 conditional

burn-flex has significantly fewer dependencies, with no dependency on ndarray itself, no macro utility crates, and no BLAS bindings.

13. Codebase Size

burn-flex has roughly 2x the code. This is because:

1. Flex implements all ops from scratch (ndarray delegates to the ndarray crate's built-in ops)
2. Flex has dedicated optimized implementations (pool, conv, reduce, cumulative, gather/scatter)
3. Flex has more comprehensive dtype handling (f16/bf16 paths for every op)
4. Flex has explicit contiguous/non-contiguous fast paths throughout

14. Testing

15. Performance Summary

All benchmarks on Apple M3 Max, default features enabled.

Compute Performance

Genuine algorithmic and library improvements:

Structural Improvements

These reflect changes in how operations are represented and executed, not pure compute speedups. burn-ndarray eagerly materializes data where burn-flex uses zero-copy views or separated storage.

Note on quantization: burn-ndarray simulates quantization by dequantizing to f32 for most operations. The quantized speedups reflect the difference between simulated and native execution, not equivalent algorithms running at different speeds.

Where NdArray Wins

These are specific edge cases where NdArray's ndarray-based internals have an advantage.

16. Why Replace burn-ndarray?

The ndarray crate has been slow to accept contributions and evolve. Burn's CPU backend inherits these constraints:

• No f16/bf16: Models using half-precision weights must convert to f32. An f16 PR has been open for a long time with no clear timeline.
• 6-dimension limit: Hard-coded in reshape macros, cannot be fixed without upstream changes.
• Unsigned strides: usize-only strides make zero-copy flip impossible.
• Simulated quantization: No native quantized storage; dequantize/requantize on every op.
• COW limitations: NdArrayStorage::Borrowed always returns false for is_unique(), preventing in-place mutation of externally loaded data.

burn-flex was built to address these gaps without waiting on upstream. It is not intended to compete with CubeCL CPU, which targets high-performance computation through operator fusion and just-in-time compilation via LLVM. The goal is to provide a lightweight, portable replacement for burn-ndarray that works today on platforms CubeCL CPU cannot target (no_std, WASM, embedded).

17. What burn-flex Adds

1. f16/bf16 support: Native arithmetic on half-precision types. Enables running models that use f16 weights without conversion.

2. No dimension limit: Arbitrary tensor rank (ndarray is limited to 6).

3. Zero-copy flip/unfold/expand: Signed strides enable O(1) flip. Unfold returns a strided view instead of materializing all windows.

4. Unified 3D conv/pool: Single implementation covers 1D/2D/3D, reducing code paths and potential for inconsistencies.

5. Native quantization: Stores scales separately for direct scale * x_q dequantization instead of reparsing packed bytes on every access. Zero-copy layout ops on quantized tensors.

6. Fewer dependencies: 7 required deps vs 12. No ndarray, no matrixmultiply, no paste, no const-random, no BLAS bindings.

7. Simpler type system: Flex vs NdArray<E, I, Q>. No generic parameters, no element trait hierarchy (FloatNdArrayElement, IntNdArrayElement, NdArrayElement, ExpElement).

8. Real FFT: Forward (rfft) and inverse (irfft) real FFT with complex packing, SIMD butterflies, and compile-time twiddle tables. Works in no_std (rustfft/realfft require std). NdArray implements neither.

18. What burn-ndarray Has That burn-flex Does Not

1. BLAS acceleration: Feature flags for Accelerate (macOS), OpenBLAS, and Netlib BLAS. These can outperform gemm for very large matmuls on specific hardware. burn-flex relies solely on the gemm crate.

2. SIMD conv/pool kernels: burn-ndarray has dedicated macerator-based SIMD kernels for convolution and pooling. burn-flex delegates to gemm's SIMD.

3. export_tests feature: burn-ndarray serves as a reference implementation for some burn-cubecl kernels via export_tests.

19. Migration Path

For Burn users switching from burn-ndarray to burn-flex:

20. Conclusion

burn-flex is a from-scratch replacement for burn-ndarray, motivated by ndarray's lack of f16/bf16 support, 6-dimension limit, simulated quantization, and slow pace of upstream development. It implements all required Backend traits (FloatTensorOps, IntTensorOps, BoolTensorOps, QTensorOps, ModuleOps, ActivationOps, TransactionOps) and passes the same test suite.

Performance gains come in two forms: compute improvements (1.1-9.7x) from better libraries and algorithms, and structural improvements (up to 166,000x) from representing operations as zero-copy views instead of eagerly materializing data. Memory usage is significantly reduced through Arc-based COW and in-place mutation.

The only capabilities lost are optional BLAS acceleration (replaced by the gemm crate, which is faster in most benchmarks) and the export_tests reference implementation feature.