Crate burn_tensor

Expand description

Burn’s tensor API.

Tensor<D, K> is a tensor of rank D and kind K (Float by default, Int or Bool). Its Device decides which backend runs it, at runtime: the same Tensor<2> can live on CUDA, wgpu or the CPU, and code that uses tensors has no backend type parameter. Applications use this crate through burn::tensor.

use burn_tensor::{Device, Tensor, s};

let device = Device::default();
let x = Tensor::<2>::from_floats([[1.0, 2.0], [3.0, 4.0]], &device);
let y = x.clone().matmul(x.transpose()).exp();
let first_row = y.slice(s![0..1, ..]);

Device: backend selection, configuration and enumeration. A backend’s constructor (Device::cuda, Device::wgpu, Device::flex, …) exists when its feature is enabled.
TensorData, DType and Shape: tensor contents and metadata, independent of any backend.
activation, loss and module: functional forms of activations, losses and neural network operations such as convolution and pooling.
quantization, grid and distributed: quantized tensors, grid sampling and collective operations across devices.
einsum!, assert_shape! and s!: macros for Einstein summation, shape checks and slicing.

With the autodiff feature, create tensors on a device.autodiff() device and mark source leaves with Tensor::require_grad to record gradients.

§Feature flags

std (default): standard library support. Without it the crate is no_std with alloc.
Backends: cuda, rocm, wgpu, metal, vulkan, webgpu, cpu, flex, and the deprecated ndarray and tch.
autodiff, fusion: backend decorators.
remote, remote-server, remote-websocket: devices hosted by another machine.
capture: record operation graphs instead of executing them.
extension: access to backend primitives for backend extensions.
tracing: instrument operations with the tracing crate.

§Note for contributors: `*_impl` helpers

Throughout this crate (e.g. in tensor::api::float, tensor::api::int, tensor::api::bool, tensor::api::cast, and tensor::activation), public generic methods on Tensor<D, K> that need to call into burn_dispatch are routed through small non-generic helper functions named *_impl, grouped together at the bottom of each file under a banner like:

// =====================================================================
// Non-generic implementation helpers (outlined from the generic API).
// =====================================================================

These helpers take and return only BridgeTensor (a type-erased blob — no DispatchTensor or other burn_dispatch types appear in their signatures). Because the helpers are not generic over D, they are compiled once, and the MIR of the public generic methods does not mention any burn_dispatch types. Downstream crates that monomorphize the public generic API therefore never have to resolve the cubecl-backed type tree, which drastically cuts compile times for user code.

When adding a new public method that calls a Dispatch::* op, follow the same pattern: keep the generic method body thin and forward to a non-generic *_impl helper alongside the existing ones.

Modules§

activation: The activation module.
cast: Tensor element casting.
container: The container module.
distributedstd: Distributed execution utilities.
grid: The grid module.
kind: The tensor kind module.
loss: The loss module.
module: The neural network module.
ops: Operations on tensors module.
quantization: Tensor quantization module.

Macros§

assert_shape: Asserts a tensor’s rank, at compile time unless the pattern has .., and its axis sizes at runtime.
debug_assert_shape: Same as assert_shape!, but compiled out unless debug_assertions is on.
einsum: Computes Einstein summation from a literal equation.
s: Slice argument constructor for tensor indexing.
shape

Structs§

BackTrace: Contains the backtrace information if available.
Bool: A type-level representation of the kind of a bool tensor.
Bytes: A buffer similar to Box<[u8]> that supports custom memory alignment and allows trailing uninitialized bytes.
Device: A high-level device handle for tensor operations.
DeviceConfig: Configuration options used to initialize a device.
DeviceFilter: A set of DeviceTypes passed to Device::enumerate.
DeviceId: The device id.
DeviceSettings: Settings controlling the default data types for a specific device.
Devices: A collection of Devices returned by Device::enumerate.
DimIter: Iterator given by (Tensor::iter_dim).
DistributionSampler: Distribution sampler for random value of a tensor.
EinsumOperand: An einsum operand whose rank is determined at runtime.
Float: A type-level representation of the kind of a float tensor
Graph: A captured computation.
IndexWrap: Wraps an index with negative indexing support.
Int: A type-level representation of the kind of a int tensor.
MemoryPoolUsage: One reading of a device allocator’s state, across the runtime’s streams and pools.
PrintOptions: Options for Tensor pretty printing
ProfileDuration: Result from profiling between two measurements. This can either be a duration or a future that resolves to a duration.
ProfileOptions: How a profiled window treats the work a backend still holds in a queue when the window closes.
ProfileTicks: Start and end point for a profile. Can be turned into a duration.
Shape: Reexport to avoid annoying rust-analyzer bug where it imports the module instead of the macro Shape of a tensor.
Slice: A slice specification for a single tensor dimension.
SliceIter: Defines an Iterator over a Slice.
SlicedPoolReport: One dynamic pool’s measured state, in the order allocations are routed through the pools.
SmallVec: Reexport for use in macros A Vec-like container that can store a small number of elements inline.
StreamId: Unique identifier representing the stream on which work is submitted.
Tensor: A tensor with a given backend, shape and data type.
TensorCreationOptions: Options for tensor creation.
TensorData: Data structure for tensors.
Tolerance: The tolerance used to compare to floating point numbers.
Transaction: A transaction can read multiple tensors at once with a single operation improving compute utilization with optimized laziness.
bf16: A 16-bit floating point type implementing the bfloat16 format.
f16: A 16-bit floating point type implementing the IEEE 754-2008 standard binary16 a.k.a “half” format.
flex32: A floating point type with relaxed precision, minimum f16, max f32.

Enums§

AllocationProperty: The kind of allocation behind the Bytes type.
BoolStore: Data type used to store boolean values.
DType
DataError: Errors that can occur while accessing or converting TensorData.
DeviceError: Errors returned by device-related operations.
DeviceIndex: Selector for the hardware index of a backend whose devices are simply indexed (e.g. CUDA, ROCm).
DeviceKind: Selector for the more flexible backends whose device handle is a tagged enum (e.g. WGPU, which can target a discrete/integrated/virtual GPU, a CPU adapter, an externally-created wgpu setup, or just “best available”).
DeviceType: Represents the devices that can be used.
Distribution: Distribution for random value of a tensor.
DistributionSamplerKind: Distribution sampler kind for random value of a tensor.
ExecutionError: An error that can happen when syncing a device.
ExpressionError: Common Expression Error.
FloatDType
IndexingUpdateOp: Computation to be used to update the existing values in indexed assignment operations (scatter/select).
IntDType
MetadataError: Reexport to avoid annoying rust-analyzer bug where it imports the module instead of the macro Error that can occur when attempting to modify shapes.
Scalar: A scalar element.
TensorReadError: Errors that can occur while reading host data from a tensor.
TimingMethod: How a benchmark’s execution times are measured.

Constants§

DEFAULT_ATOL: Default ATOL value for is_close and all_close.
DEFAULT_RTOL: Default RTOL value for is_close and all_close.

Traits§

AsIndex: Helper trait for implementing indexing with support for negative indices.
AsSize: Types which can be converted to a usize Size.
BroadcastArgs: Trait used for broadcast arguments.
Cast: Trait for types that represent a valid cast target from a tensor of kind K.
Element: Core element trait for tensor values.
ElementAdd: Element addition trait.
ElementComparison: Element ordering trait.
ElementConversion: Element conversion trait for tensor.
ElementEq: Element trait for equality of a tensor.
ElementLimits: Element limits trait.
ElementOrdered: Ordered element trait for tensor values.
ElementRandom: Element trait for random value of a tensor.
IntoPadding: Trait for types that can be used as padding specifications.
MovedimArgs: Trait used for movedim arguments
ReshapeArgs: Trait used for reshape arguments.
SliceArg: Trait for slice arguments that can be converted into an array of slices. This allows the slice method to accept both single slices (from s![..]) and arrays of slices (from s![.., ..] or [0..5, 1..3]).
SliceOps: Slice-related ops on Shape

Functions§

calculate_matmul_output: Reexport to avoid annoying rust-analyzer bug where it imports the module instead of the macro Compute the output shape for matrix multiplication with broadcasting support.
capture: Capture closure for repeated replay (see Graph).
cartesian_grid: Generates a cartesian grid for the given tensor shape on the specified device. The generated tensor is of dimension D2 = D + 1, where each element at dimension D contains the cartesian grid coordinates for that element.
check_closenessstd: Checks the closeness of two tensors and prints the results.
ravel_index: Compute the ravel index for the given coordinates.
read_sync: Read a future synchronously.
set_print_options: Set print options
try_read_sync: Read a future synchronously.
wrap_index: Wraps a dimension index to be within the bounds of the dimension size.

Type Aliases§

BoolDType: Boolean dtype.