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

Modules

activation

The activation module.

backend

The backend module.

cast

Tensor element casting.

container

The container module.

grid

The grid module.

linalg

The linalg module.

loss

The loss module.

module

The neural network module.

ops

Operations on tensors module.

quantization

Tensor quantization module.

Macros

s

Slice argument constructor for tensor indexing.

Structs

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.

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

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.

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.

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.

TransactionPrimitive

Contains all tensor primitives that are going to be read.

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.

BoundsError

Access Bounds Error.

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

IndexingUpdateOp

Computation to be used to update the existing values in indexed assignment operations (scatter/select).

IntDType

ShapeError

Error that can occur when attempting to modify shapes.

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

AsIndex

Helper trait for implementing indexing with support for negative indices.

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.

ElementRandom

Element trait for random value of a tensor.

MovedimArgs

Trait used for movedim arguments

Numeric

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

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]).

TensorKind

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

TensorMetadata

Tensor metadata trait for tensor primitive.

Functions

calculate_matmul_output

Compute the output shape for matrix multiplication with broadcasting support.

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.

expect_wrap

Wraps an index with negative indexing support.

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.

try_wrap

Wraps an index with negative indexing support.

wrap_index

Wraps a dimension index to be within the bounds of the dimension size.

Type Aliases

Device

Device type used by the backend.