Module tensor 

Tensor types and compatibility re-exports.

Modules

activation

The activation module.

cast

Tensor element casting.

container

The container module.

distributed

Distributed execution utilities.

grid

The grid module.

kind

The tensor kind module.

linalg

Compatibility path for linear algebra operations.

loss

The loss module.

module

The neural network module.

ops

Operations on tensors module.

quantization

Tensor quantization module.

signal

Compatibility path for signal processing operations.

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

DispatchTensor

A tensor that can dispatch operations to any enabled backend at runtime.

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

Gradients

Gradients container used during the backward pass.

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

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

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

GradientCheckpointingStrategy

Gradient checkpointing strategy for autodiff.

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.

PrimitiveConversionError

Error returned when a DispatchTensor cannot be converted to the requested concrete primitive type.

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.

BackendPrimitive

Trait to safely extract and wrap backend-specific primitives based on the high-level tensor kind.

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

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_closeness

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.