Crate burn_tensor

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 support auto-differentiation thanks to the AutodiffBackend trait.

Modules

activation

The activation module.

backend

The backend module.

cast

Tensor element casting.

container

The container module.

loss

The loss module.

module

The burn module.

ops

Operations on tensors module.

quantization

Tensor quantization module.

Macros

dequant_op_quant

Automatically applies dequantization -> float operation -> quantization.

make_element

Macro to implement the element trait for a type.

s

Creates a slice specification for tensor indexing operations.

Structs

Bool

A type-level representation of the kind of a bool tensor.

Bytes

A sort of Box<[u8]> that remembers the original alignment and can contain trailing uninitialized bytes.

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

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 (range).

T

Transpose marker (zero-size type). Used to sugar the transpose of a tensor, e.g.

Tensor

A tensor with a given backend, shape and data type.

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

DType

DataError

The things that can go wrong when manipulating tensor data.

Distribution

Distribution for random value of a tensor.

DistributionSamplerKind

Distribution sampler kind for random value of a tensor.

FloatDType

Precision

Element precision trait for tensor.

TensorPrimitive

A primitive tensor representation.

Constants

DEFAULT_ATOL

Default ATOL value for is_close and all_close.

DEFAULT_RTOL

Default RTOL value for is_close and all_close.

Traits

BasicAutodiffOps

Trait that list all operations that can be applied on all tensors on an autodiff backend.

BasicOps

Trait that list all operations that can be applied on all tensors.

BroadcastArgs

Trait used for broadcast arguments.

Element

Element trait for tensor.

ElementComparison

Element ordering trait.

ElementConversion

Element conversion trait for tensor.

ElementLimits

Element ordering trait.

ElementPrecision

Element precision trait for tensor.

ElementRandom

Element trait for random value of a tensor.

IndexConversion

A helper trait to convert difference indices type to a slice index.

MovedimArgs

Trait used for movedim arguments

Numeric

Trait that list all operations that can be applied on all numerical tensors.

RangesArg

Trait used for slice arguments

ReshapeArgs

Trait used for reshape arguments.

TensorKind

A type-level representation of the kind of a tensor. Metadata access is lazy.

TensorMetadata

Tensor metadata trait for tensor primitive.

Functions

argsort

Returns the indices that sort the elements of the input tensor along a given dimension.

argwhere_data

Compute the indices of the elements that are non-zero, grouped by element.

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.

chunk

Split the tensor along the given dimension into chunks.

narrow

Returns a new tensor with the given dimension narrowed to the given range.

read_sync

Read a future synchronously.

set_print_options

Set print options

sort

Sort the elements of the input tensor by value along a given dimension.

sort_with_indices

Sort the elements of the input tensor by value along a given dimension.

split

Splits the tensor along the given dimension into equally sized chunks (if possible) with size split_size. Last chunk will be smaller if the tensor size along the given dimension dim is not divisible by split_size.

split_with_sizes

Splits the tensor along the given dimension into chunks with sizes in dim according to split_sizes.

try_read_sync

Read a future synchronously.

Type Aliases

Device

Device type used by the backend.