Rust Burn Framework: Tensor Manipulation, Gradient Tracking, and Memory Allocation Strategies
This document defines a structured implementation pattern for tensor manipulation, gradient tracking, and allocation planning in Rust applications using the Burn framework. The goal is to separate mathematical tensor semantics from backend-specific execution details while preserving strong type guarantees. Burn-style tensor code typically binds tensor values to a backend, device, dimensionality, and scalar representation. This design allows numerical programs to remain portable across CPU, GPU, and autodiff-enabled execution contexts.
Let a dense tensor be represented as:
\[ T \in \mathbb{R}^{d_1 \times d_2 \times \cdots \times d_n} \]
The logical number of elements is:
\[ |T|=\prod_{i=1}^{n}d_i \]
If each scalar occupies \(b\) bytes, the minimum contiguous payload size is:
\[ M_{\text{payload}}=b\prod_{i=1}^{n}d_i \]
This value is not the full runtime allocation cost. Runtime memory may also include backend buffer alignment, cached allocations, graph metadata, temporary tensors, gradient storage, and device transfer staging buffers. For training workloads, a practical memory estimate must include parameters, activations, gradients, and temporary intermediate values:
\[ M_{\text{peak}}\approx M_{\text{parameters}}+ M_{\text{activations}}+ M_{\text{gradients}}+ M_{\text{temporaries}} \]
Gradient tracking should be enabled only for tensors that participate in optimization. Labels, masks, constants, and inference-only inputs should not be differentiated unless the mathematical objective requires it. For a scalar loss \(L(\theta)\), the training process requires the gradient vector:
\[ \nabla_{\theta}L= \left( \frac{\partial L}{\partial \theta_1}, \frac{\partial L}{\partial \theta_2}, \ldots, \frac{\partial L}{\partial \theta_k} \right) \]
Reference Implementation
use burn::prelude::*;
use burn::tensor::{Data, Tensor};
use burn::backend::Autodiff;
#[derive(Module, Debug)]
pub struct DenseBlock<B: Backend> {
weight: Param<Tensor<B, 2>>,
bias: Param<Tensor<B, 1>>,
}
impl<B: Backend> DenseBlock<B> {
pub fn new(input_dim: usize, output_dim: usize, device: &B::Device) -> Self {
let weight_data = Data::random(
[input_dim, output_dim],
burn::tensor::Distribution::Normal(0.0, 0.02),
);
let bias_data = Data::zeros([output_dim]);
Self {
weight: Param::from_tensor(Tensor::from_data(weight_data, device)),
bias: Param::from_tensor(Tensor::from_data(bias_data, device)),
}
}
pub fn forward(&self, input: Tensor<B, 2>) -> Tensor<B, 2> {
input.matmul(self.weight.val()).add(self.bias.val().unsqueeze())
}
}
pub fn tensor_payload_bytes(shape: &[usize], scalar_bytes: usize) -> usize {
shape.iter().product::<usize>() * scalar_bytes
}
pub fn normalize_batch<B: Backend>(x: Tensor<B, 2>) -> Tensor<B, 2> {
let mean = x.clone().mean_dim(0);
let centered = x - mean.unsqueeze();
let variance = centered.clone().powf_scalar(2.0).mean_dim(0);
let std = variance.add_scalar(1e-6).sqrt();
centered / std.unsqueeze()
}
pub fn training_step<B>(
model: &DenseBlock<Autodiff<B>>,
input: Tensor<Autodiff<B>, 2>,
target: Tensor<Autodiff<B>, 2>,
) -> Tensor<Autodiff<B>, 1>
where
B: Backend,
{
let prediction = model.forward(input);
let error = prediction - target;
let loss = error.powf_scalar(2.0).mean();
loss.backward()
}
fn main() {
type InnerBackend = burn_ndarray::NdArray<f32>;
type ADBackend = Autodiff<InnerBackend>;
let device = Default::default();
let model = DenseBlock::<ADBackend>::new(8, 4, &device);
let input = Tensor::<ADBackend, 2>::ones([16, 8], &device);
let target = Tensor::<ADBackend, 2>::zeros([16, 4], &device);
let payload = tensor_payload_bytes(&[16, 8], std::mem::size_of::<f32>());
println!("logical payload bytes: {}", payload);
let gradients = training_step(&model, input, target);
println!("backward result: {:?}", gradients);
}The dense block accepts a two-dimensional input tensor and applies affine transformation. For input \(X\), weight matrix \(W\), and bias vector \(b\), the operation is:
\[ Y=XW+b \]
If \(X\in\mathbb{R}^{B\times d_{\text{in}}}\), \(W\in\mathbb{R}^{d_{\text{in}}\times d_{\text{out}}}\), and \(b\in\mathbb{R}^{d_{\text{out}}}\), then:
\[ Y\in\mathbb{R}^{B\times d_{\text{out}}} \]
The batch dimension \(B\) may vary between training steps, but the invariant dimensions \(d_{\text{in}}\) and \(d_{\text{out}}\) must remain consistent with the model parameters. This distinction should be part of the public API contract.
Allocation strategy is especially important when the same model must run under different backends. A CPU backend may tolerate additional temporary allocations, while a GPU backend may require stricter control over peak memory. The maximum safe batch size can be approximated as:
\[ B_{\max}= \left\lfloor \frac{M_{\text{available}}-M_{\text{fixed}}} {M_{\text{per-sample}}} \right\rfloor \]
A robust implementation should expose backend type, device selection, scalar precision, batch size, and gradient mode as explicit configuration parameters. Model code should not silently choose execution devices or enable autodiff for values that do not require derivatives. This separation improves reproducibility, reduces memory waste, and makes the training pipeline easier to audit.