Tensor

As previously explained in the model section, Tensor has a const generic for its dimensionality D and an optional tensor kind:

Tensor<D>           // Float tensor (default)
Tensor<D, Float>    // Explicit float tensor
Tensor<D, Int>      // Int tensor
Tensor<D, Bool>     // Bool tensor

Data types

A tensor has an element type, but unlike its Float/Int/Bool kind that type is a runtime property rather than a generic parameter on Tensor. When a tensor is created without an explicit dtype (for example with Tensor::zeros, Tensor::ones or from_floats), the element type comes from a set of default data types that each device keeps. You can inspect the current defaults with Device::settings:

use burn::tensor::Device;

let device = Device::wgpu(Default::default());
let settings = device.settings();
// settings.float_dtype, settings.int_dtype, settings.bool_dtype

To use a different default (e.g. f16 for floats), configure the device with Device::configure before creating any tensor on it:

use burn::tensor::{Device, FloatDType, IntDType};

let mut device = Device::wgpu(Default::default());
device.configure((FloatDType::F16, IntDType::I32))?;

// Float tensors created after this default to F16
let floats = Tensor::<2>::zeros([2, 3], &device);

Default data types lock on first use. A device's defaults can only be initialized once, and that initialization must happen before the first tensor operation on the device. The first tensor created on an unconfigured device permanently locks the defaults to the backend's values. Any later call to configure (or set_default_dtypes) for that device returns a DeviceError::AlreadyInitialized error instead of silently changing (or ignoring) the width.

Because the default width is a per-device property, you cannot have two different default float widths active on the same device within a single process. If you need values at more than one precision, create the other tensors with an explicit dtype by passing a (&device, dtype) tuple as the creation options (this tuple is just a convenient conversion into TensorCreationOptions):

use burn::tensor::DType;

// device defaults to f32
let x = Tensor::<2>::zeros([2, 3], &device);                   // f32
let x_f64 = Tensor::<2>::zeros([2, 3], (&device, DType::F64)); // explicit f64

To convert an existing tensor to another element type, use cast:

let x_f64 = x.cast(FloatDType::F64); // convert the f32 tensor above to f64

Burn Tensors are defined by the number of dimensions D in their declaration as opposed to their shape. The actual shape of the tensor is inferred from its initialization. For example, a Tensor of size (5,) is initialized as below:

let floats = [1.0, 2.0, 3.0, 4.0, 5.0];

// Get the default device
let device = Default::default();

// correct: Tensor is 1-Dimensional with 5 elements
let tensor_1 = Tensor::<1>::from_floats(floats, &device);

// incorrect: let tensor_1 = Tensor::<5>::from_floats(floats, &device);
// this will lead to an error and is for creating a 5-D tensor

Initialization

Burn Tensors are primarily initialized using the from_data() method which takes the TensorData struct as input. The TensorData fields are private: shape() and dtype() return the shape and data type, and the values, stored as bytes, can be accessed via any of the following methods: as_slice, as_mut_slice, try_to_vec, try_to_vec_as and iter. To retrieve the data from a tensor, the method .to_data() should be employed when intending to reuse the tensor afterward. Alternatively, .into_data() is recommended for one-time use. Let's look at a couple of examples for initializing a tensor from different inputs.

// Initialization on the selected runtime device
let tensor_1 = Tensor::<1>::from_data([1.0, 2.0, 3.0], &device);

// Initialization from TensorData
let tensor_2 = Tensor::<1>::from_data(TensorData::from([1.0, 2.0, 3.0]), &device);

// Initialization using from_floats (Recommended for f32 ElementType)
// Will be converted to TensorData internally.
let tensor_3 = Tensor::<1>::from_floats([1.0, 2.0, 3.0], &device);

// Initialization of Int Tensor from array slices
let arr: [i32; 6] = [1, 2, 3, 4, 5, 6];
let tensor_4 = Tensor::<1, Int>::from_data(TensorData::from(&arr[0..3]), &device);

// Initialization from a custom type

struct BodyMetrics {
    age: i8,
    height: i16,
    weight: f32
}

let bmi = BodyMetrics{
        age: 25,
        height: 180,
        weight: 80.0
    };
let data  = TensorData::from([bmi.age as f32, bmi.height as f32, bmi.weight]);
let tensor_5 = Tensor::<1>::from_data(data, &device);

Reading a tensor's data is also where a failed computation surfaces. into_data() and to_data() panic in that case; use try_into_data() or try_to_data() to handle the ExecutionError instead.

Ownership and Cloning

Almost all Burn operations take ownership of the input tensors. Therefore, reusing a tensor multiple times will necessitate cloning it. Let's look at an example to understand the ownership rules and cloning better. Suppose we want to do a simple min-max normalization of an input tensor.

let input = Tensor::<1>::from_floats([1.0, 2.0, 3.0, 4.0], &device);
let min = input.min();
let max = input.max();
let input = (input - min).div(max - min);

With PyTorch tensors, the above code would work as expected. However, Rust's strict ownership rules will give an error and prevent using the input tensor after the first .min() operation. The ownership of the input tensor is transferred to the variable min and the input tensor is no longer available for further operations. Burn Tensors like most complex primitives do not implement the Copy trait and therefore have to be cloned explicitly. Now let's rewrite a working example of doing min-max normalization with cloning.

let input = Tensor::<1>::from_floats([1.0, 2.0, 3.0, 4.0], &device);
let min = input.clone().min();
let max = input.clone().max();
let input = (input.clone() - min.clone()).div(max - min);
println!("{}", input.to_data());// Success: [0.0, 0.33333334, 0.6666667, 1.0]

// Notice that max, min have been moved in last operation so
// the below print will give an error.
// If we want to use them for further operations,
// they will need to be cloned in similar fashion.
// println!("{:?}", min.to_data());

We don't need to be worried about memory overhead because with cloning, the tensor's buffer isn't copied, and only a reference to it is increased. This makes it possible to determine exactly how many times a tensor is used, which is very convenient for reusing tensor buffers or even fusing operations into a single kernel (burn-fusion). For that reason, we don't provide explicit inplace operations. If a tensor is used only one time, inplace operations will always be used when available.

Shape Assertions

The rank of a tensor is checked by the type system, and let [b, t, c] = x.dims(); names its axes. To check the size of each axis, use the shape macros from the prelude:

let [batch_size, seq_length, _] = x.dims();
assert_shape!(x, [_, _, 80]);
assert_shape!(y, [batch_size, seq_length, 256]);
assert_shape!(flat, [batch_size * seq_length, 256]);
assert_shape!(w, [.., 256]); // any rank, last axis checked
debug_assert_shape!(z, [batch_size, _, 256]);

See Forward Contract for the slot syntax and guidance on which macro to use where.

Tensor Operations

Normally with PyTorch, explicit inplace operations aren't supported during the backward pass, making them useful only for data preprocessing or inference-only model implementations. With Burn, you can focus more on what the model should do, rather than on how to do it. We take the responsibility of making your code run as fast as possible during training as well as inference. The same principles apply to broadcasting; all operations support broadcasting unless specified otherwise.

Here, we provide a list of all supported operations along with their PyTorch equivalents. Note that for the sake of simplicity, we ignore type signatures. For more details, refer to the full documentation.

Basic Operations

Those operations are available for all tensor kinds: Int, Float, and Bool.

Numeric Operations

Those operations are available for numeric tensor kinds: Float and Int.

Float Operations

Those operations are only available for Float tensors.

Int Operations

Those operations are only available for Int tensors.

Bool Operations

Those operations are only available for Bool tensors.

Quantization Operations

Those operations are only available for Float tensors on backends that implement quantization strategies.

Einstein Summation

Use einsum! to express tensor contractions with an equation. Each letter names an input axis; commas separate operands, and -> lists the output axes in order. Values are multiplied along matching labels and summed over labels omitted from the output. For example, "ij,jk->ik" multiplies two matrices and sums over j.

use burn::tensor::{Tensor, einsum};

// Queries: [batch, queries, channels]; features: [batch, channels, height, width].
// Sum over channels to produce [batch, queries, height, width].
let masks = einsum!("bqc,bchw->bqhw", &queries, &features);

// Runtime equations accept operands of different ranks through `.into()`.
let equation = "ij,j->i";
let output = Tensor::<1>::einsum(equation, [matrix.into(), vector.into()]);

The macro checks the equation, operand count, and statically determined ranks at compile time. Shapes and broadcasting are checked at runtime.

Broadcasting and diagonals

Matching labels across operands must have equal sizes or a size of one, which broadcasts to the other size. Repeated labels within a single operand extract a diagonal and require equal axis sizes; singleton broadcasting does not apply there.

An ellipsis (...) matches zero or more axes. Ellipsis dimensions broadcast from the right: "...ij,...jk->...ik" can multiply shapes [2, 3, 4] and [4, 5] to produce [2, 3, 5]. Supply the result type when an output ellipsis has unknown width:

let result: Tensor<3> = einsum!("...ij,...jk->...ik", batches, matrix);

Omitting ... from an explicit output sums over its dimensions. With no ->, the output contains the ellipsis first, followed by labels occurring exactly once across all inputs, sorted A-Z, then a-z. Each axis label is a single letter from a-z or A-Z. Uppercase and lowercase letters are distinct labels.

Types and limitations

Float and Int operands must share their kind, dtype, and device. Scalar results have shape [1]; an empty input subscript also accepts a tensor of shape [1]. Quantized operands are unsupported. Operands are contracted from left to right using existing tensor operations, so floating-point contractions support autodiff. The implementation does not search for an optimized contraction order.

Run the matrix multiplication, mask prediction, and gradient demo from the repository:

cargo run -p burn-tensor --example einsum --features flex,autodiff

Activation Functions

Grid Functions

Linalg Functions

Signal Processing Functions

Signal-processing helpers live in the burn-signal extension crate and operate on real-valued float tensors. Enable Burn's optional signal feature to use burn::signal; the compatibility path burn::tensor::signal exports the same functions. The direct burn_tensor::signal and burn_core::tensor::signal paths have been removed.

FFT operations use the SignalOps backend extension. Windows and STFT/ISTFT compose tensor operations. Enable autodiff to differentiate through FFTs. Remote servers and captured-graph interpreters must register burn_signal::register_fft_ops in their custom-operation registry (with the router feature enabled).

FFT length n (and n_fft in STFT) must currently be a power of two: when n is Some(size), the input is truncated or zero-padded to size and the output has size / 2 + 1 frequency bins. Non-power-of-two sizes panic at the public API boundary; general arbitrary-size DFT support (Bluestein's algorithm) is a tracked follow-up.

stft and istft share a StftOptions struct with fields n_fft, hop_length, win_length, center, and onesided. Use StftOptions::new(n_fft) for PyTorch-style defaults (hop_length = n_fft / 4, win_length = None, center = true, onesided = true). The option set is validated on entry to both stft and istft; n_fft must be a power of two and hop_length <= effective_win_length (the COLA prerequisite for invertibility).

Displaying Tensor Details

Burn provides flexible options for displaying tensor information, allowing you to control the level of detail and formatting to suit your needs.

Basic Display

To display a detailed view of a tensor, you can simply use Rust's println! or format! macros:

let tensor = Tensor::<2>::full([2, 3], 0.123456789, &Default::default());
println!("{}", tensor);

This will output:

Tensor {
  data:
[[0.12345679, 0.12345679, 0.12345679],
 [0.12345679, 0.12345679, 0.12345679]],
  shape:  [2, 3],
  device:  Cpu,
  backend:  "flex",
  kind:  "Float",
  dtype:  "f32",
}

Controlling Precision

You can control the number of decimal places displayed using Rust's formatting syntax:

#![allow(unused)]
fn main() {
println!("{:.2}", tensor);
}

Output:

Tensor {
  data:
[[0.12, 0.12, 0.12],
 [0.12, 0.12, 0.12]],
  shape:  [2, 3],
  device:  Cpu,
  backend:  "flex",
  kind:  "Float",
  dtype:  "f32",
}

Global Print Options

For more fine-grained control over tensor printing, Burn provides a PrintOptions struct and a set_print_options function:

use burn::tensor::{set_print_options, PrintOptions};

let print_options = PrintOptions {
    precision: Some(2),
    ..Default::default()
};

set_print_options(print_options);

Options:

• precision: Number of decimal places for floating-point numbers (default: None)

• threshold: Maximum number of elements to display before summarizing (default: 1000)

• edge_items: Number of items to show at the beginning and end of each dimension when summarizing (default: 3)

  Checking Tensor Closeness

  Burn provides a utility function check_closeness to compare two tensors and assess their similarity. This function is particularly useful for debugging and validating tensor operations, especially when working with floating-point arithmetic where small numerical differences can accumulate. It's also valuable when comparing model outputs during the process of importing models from other frameworks, helping to ensure that the imported model produces results consistent with the original.

  Here's an example of how to use check_closeness:

  use burn::tensor::{check_closeness, Device, Tensor};
  
  let device = Device::flex();
  let tensor1 = Tensor::<1>::from_floats(
      [1.0, 2.0, 3.0, 4.0, 5.0, 6.001, 7.002, 8.003, 9.004, 10.1],
      &device,
  );
  let tensor2 = Tensor::<1>::from_floats(
      [1.0, 2.0, 3.0, 4.000, 5.0, 6.0, 7.001, 8.002, 9.003, 10.004],
      &device,
  );
  
  check_closeness(&tensor1, &tensor2);

  The check_closeness function compares the two input tensors element-wise, checking their absolute differences against a range of epsilon values. It then prints a detailed report showing the percentage of elements that are within each tolerance level.

  The output provides a breakdown for different epsilon values, allowing you to assess the closeness of the tensors at various precision levels. This is particularly helpful when dealing with operations that may introduce small numerical discrepancies.

  The function uses color-coded output to highlight the results:

  • Green [PASS]: All elements are within the specified tolerance.
  • Yellow [WARN]: Most elements (90% or more) are within tolerance.
  • Red [FAIL]: Significant differences are detected.

  This utility can be invaluable when implementing or debugging tensor operations, especially those involving complex mathematical computations or when porting algorithms from other frameworks. It's also an essential tool when verifying the accuracy of imported models, ensuring that the Burn implementation produces results that closely match those of the original model.