Crate burn_tensor

Expand description

This library provides the core abstractions required to run tensor operations with Burn. Tensors are generic over the backend to allow users to perform operations using different Backend implementations. Burn’s tensors also support auto-differentiation thanks to the AutodiffBackend trait.

§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.
linalg: The linalg module.
loss: The loss module.
module: The neural network module.
ops: Operations on tensors module.
quantization: Tensor quantization module.
signal: The signal processing module.

Macros§

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.
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.
PrintOptions: Options for Tensor pretty printing
Shape: Shape of a tensor.
Slice: A slice specification for a single tensor dimension.
SliceIter: Defines an Iterator over a Slice.
SmallVec: A Vec-like container that can store a small number of elements inline.
StreamId: Unique identifier that can represent a stream based on the current thread id.
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.

Enums§

AllocationProperty: The kind of allocation behind the Bytes type.
BoolStore: Data type used to store boolean values.
DType
DataError: The things that can go wrong when manipulating tensor data.
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: Error that can occur when attempting to modify shapes.
Scalar: A scalar element.

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.
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: 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.