Part 1: Tensors

As you begin your implementation, the first thing you will encounter is the tensor. In order to know how deep learning works, you must understand this fundamental structure.

A tensor is simply a multidimensional array. You can think of it as a generalization of a matrix. In deep learning, all data is represented as a tensor. But why this specific structure? One of their key advantages is that they can perform mathematical operations efficiently and in parallel on GPUs, making large-scale computation possible.

In this section, we will explore how tensors are created and utilized with the Burn crate. By examining some tensor operations, you will gain a deeper understanding of tensors.

Project setup

Let’s start by setting up a small Rust project using the Burn deep learning framework.

Create a new binary project:

cargo new simple-regression --bin

Then, open the Cargo.toml file and add Burn as a dependency:

[package]
name = "simple-regression"
version = "0.1.0"
edition = "2021"

[dependencies]
burn = { version = "0.17.1", features = ["ndarray"] }

Here we use the ndarray backend, which performs computations on the CPU.
Other backends like wgpu (for GPU acceleration) can be used later.

Define tensor

Let’s define a simple tensor and inspect it:

use burn::{tensor::Tensor};
use burn::backend::NdArray;

fn main() {
    fn print_tensor<const D: usize>(label: &str, t: &Tensor<B, D>) {
        let data: Vec<f32> = t.clone().into_data().convert::<f32>().to_vec().unwrap();
        println!("{label} shape={:?}, values={:?}", t.shape(), data);
    }

    type B = NdArray;

    let t = Tensor::<B, 2>::from_floats([[1.0, 2.0, 3.0],
                                            [4.0, 5.0, 6.0]],
                                        &Default::default());

    print_tensor("Original", &t);
}

• Tensor::<B, 2> — The tensor has two parameters:

  • Backend (B) — determines how the tensor is stored and computed (NdArray, Wgpu, etc.).

  • Rank (2) — the number of dimensions.

• The shape of t is [2, 3].

• print_tensor is a helper function that converts the tensor to a vector of f32 values and prints its shape and contents.

Output:

Original shape=Shape { dims: [2, 3] }, values=[1.0, 2.0, 3.0, 4.0, 5.0, 6.0]

sum_dim: Summing Over a Dimension

The sum_dim function computes the sum of all elements along a specified dimension (or axis) of the tensor, reducing that dimension by aggregation.

let sum0 = t.clone().sum_dim(0);
print_tensor("After sum_dim(0)", &sum0); // After sum_dim(0) shape=Shape { dims: [1, 3] }, values=[5.0, 7.0, 9.0]

Output:

After sum_dim(0) shape=Shape { dims: [1, 3] }, values=[5.0, 7.0, 9.0]

The tensor’s shape [2, 3] has two rows (dimension 0).
By summing along that axis, we combine the rows:

[[1 + 4, 2 + 5, 3 + 6]] = [[5, 7, 9]]

squeeze: Removing Dimensions

squeeze function removes the specified dimension from the tensor if its size is one, effectively reducing the tensor’s number of dimensions by one.

let sum0_squeezed = sum0.squeeze::<1>(0);
print_tensor("After squeeze::<0>()", &sum0_squeezed); // After squeeze::<0>() shape=Shape { dims: [3] }, values=[5.0, 7.0, 9.0]

Output:

After squeeze::<0>() shape=Shape { dims: [3] }, values=[5.0, 7.0, 9.0]

We removed the first dimension (of size 1), changing the shape from [1, 3] → [3].
Now, the tensor’s type is Tensor::<B, 1>.

Summing Along the Second Dimension

Let’s repeat the process, but sum along the second axis (dimension 1):

let sum1 = t.sum_dim(1);
print_tensor("After sum_dim(1)", &sum1);
let sum1_squeezed = sum1.squeeze::<1>(1);
print_tensor("After squeeze::<1>()", &sum1_squeezed);

Output:

After sum_dim(1) shape=Shape { dims: [2, 1] }, values=[6.0, 15.0]
After squeeze::<1>() shape=Shape { dims: [2] }, values=[6.0, 15.0]

Here, summing across columns gives:

[1 + 2 + 3 = 6], [4 + 5 + 6 = 15]

The shape [2, 3] becomes [2, 1], and after squeezing, [2].

Squeezing a Rank-3 Tensor

Now, let’s try a tensor with three dimensions:

let tensor = Tensor::<B, 3>::from_data(
    [[[3.0, 4.9, 2.0]], [[2.0, 1.9, 3.0]], [[4.0, 5.9, 8.0]]],
    &Default::default(),
);
print_tensor("Original", &tensor);

let squeezed = tensor.squeeze::<2>(1);
print_tensor("After squeeze::<2>()", &squeezed);

Output:

Original shape=Shape { dims: [3, 1, 3] }, values=[3.0, 4.9, 2.0, 2.0, 1.9, 3.0, 4.0, 5.9, 8.0]
After squeeze::<2>() shape=Shape { dims: [3, 3] }, values=[3.0, 4.9, 2.0, 2.0, 1.9, 3.0, 4.0, 5.9, 8.0]

The shape [3, 1, 3] becomes [3, 3] because we removed the second dimension (of size 1).
The resulting tensor type is Tensor::<B, 2>.

Conclusion

Tensors are the fundamental part of the deep learning computations.

In this part, we leanred how to define data with tensor type, manipulate dimensions with sum_dim and squeeze, and track how shapes transform thorugh these operations.

In the next part, we’ll start generating our own train data.