Struct Tensor 

pub struct Tensor<const D: usize, K = Float>
where
    K: Basic,
{ /* private fields */ }

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

Indexing

Indexing a tensor can be done using slice for all tensor types or select for numeric types.

Example

use burn_tensor::Tensor;
use burn_tensor::Int;

let device = Default::default();

let tensor = Tensor::<2>::from_data(
    [
        [3.0, 4.9, 2.0],
        [2.0, 1.9, 3.0],
        [6.0, 1.5, 7.0],
        [3.0, 4.9, 9.0],
    ],
    &device,
);

// Slice the tensor to get the second and third rows:
// [[2.0, 1.9, 3.0], [6.0, 1.5, 7.0]]
// The resulting tensor will have dimensions [2, 3].
let slice = tensor.clone().slice([1..3]);
println!("{slice}");

// Slice the tensor to get the first two rows and the first 2 columns:
// [[3.0, 4.9], [2.0, 1.9]]
// The resulting tensor will have dimensions [2, 2].
let slice = tensor.clone().slice([0..2, 0..2]);
println!("{slice}");

// Index the tensor along the dimension 1 to get the elements 0 and 2:
// [[3.0, 2.0], [2.0, 3.0], [6.0, 7.0], [3.0, 9.0]]
// The resulting tensor will have dimensions [4, 2]
let indices = Tensor::<1, Int>::from_data([0, 2], &device);
let indexed = tensor.select(1, indices);
println!("{indexed}");

Implementations

impl<const D: usize> Tensor<D>

pub fn backward(&self) -> Gradients

Computes gradients by backpropagating from this tensor.

The tensor must participate in an autodiff graph. Backward consumes the recorded steps reachable from this tensor, even though this method borrows it. Clones share those steps; cloning a tensor does not preserve its backward tape.

Burn does not support retaining the graph for another backward (retain_graph). Combine losses that share intermediates into one loss before calling backward, or recompute the forward pass for each backward. Fresh forwards and existing branches can reuse parameter leaves, including leaves used by an earlier backward, as long as they do not depend on consumed intermediates. Gradient checkpointing does not change this contract.

To intentionally stop gradients through earlier operations, use detach. This starts a new lineage; it does not restore the earlier graph.

Panics

Panics if autodiff is disabled, the tensor doesn’t participate in a recorded graph, backward has already consumed this tensor’s step, or any required intermediate step has been consumed. Consumed ancestry is rejected before any additional steps are consumed, leaving fresh branches available for backward. Distributed backward also panics if a distributed parameter uses a different backend than the loss.

pub fn grad(&self, grads: &Gradients) -> Option<Tensor<D>>

Returns this tensor’s retained gradient, if present.

The returned gradient doesn’t have an autodiff association. This returns None when the tensor is outside autodiff, doesn’t retain gradients, or isn’t present in grads. Repeated calls return handles to the same gradient. If the gradient is only needed once, prefer grad_remove, which can enable in-place optimizations.

pub fn grad_remove(&self, grads: &mut Gradients) -> Option<Tensor<D>>

Removes and returns this tensor’s retained gradient, if present.

The returned gradient doesn’t have an autodiff association. This returns None when the tensor is outside autodiff, doesn’t retain gradients, or isn’t present in grads.

pub fn grad_replace(&self, grads: &mut Gradients, grad: Tensor<D>)

Replaces this tensor’s entry in grads with grad.

Panics

Panics if this tensor isn’t associated with autodiff, if grad has an autodiff association, or if the tensors use incompatible backends.

pub fn is_tracked(&self) -> bool

Returns whether this tensor’s node is marked for autodiff graph participation.

A tensor can have autodiff enabled without being tracked, such as a constant that doesn’t require gradients. is_autodiff reports whether autodiff is enabled for the tensor, while is_require_grad reports whether its gradient is retained after backward.

This reports graph participation and doesn’t indicate whether the graph tape has already been consumed.

impl<const D: usize, K> Tensor<D, K>

where K: Autodiff,

pub fn is_autodiff(&self) -> bool

Returns whether autodiff is enabled for this tensor.

This doesn’t indicate whether the tensor participates in a recorded graph or retains its gradient. Inspect those properties with Tensor::is_tracked() and is_require_grad, respectively.

pub fn inner(self) -> Tensor<D, K>

Returns this tensor without its autodiff association.

If the tensor uses autodiff, this moves it to the inner backend and drops any graph reference it carries. If autodiff is already disabled, the tensor is returned unchanged. This operation is idempotent.

This is equivalent to without_autodiff. inner reflects the underlying backend-decorator model, while without_autodiff describes the operation in terms of the high-level tensor API.

pub fn without_autodiff(self) -> Tensor<D, K>

Returns this tensor without its autodiff association.

If the tensor uses autodiff, this moves it to the inner backend and drops any graph reference it carries. If autodiff is already disabled, the tensor is returned unchanged. This operation is idempotent.

Unlike detach, which severs the tensor from its current graph but keeps its autodiff association, the returned tensor pays no autodiff dispatch for subsequent operations involving only tensors without autodiff. When combined with a tensor that still uses autodiff, it is treated as a constant for that operation.

pub fn autodiff(self) -> Tensor<D, K>

Returns this tensor with autodiff enabled.

If autodiff is disabled, this associates the tensor with the autodiff backend using the default gradient-checkpointing strategy. If autodiff is already enabled, the tensor and its current strategy are returned unchanged. This operation is idempotent.

Enabling autodiff does not make a floating-point tensor require gradients. Use require_grad when its gradient should be retained during the backward pass.

pub fn from_inner(inner: Tensor<D, K>) -> Tensor<D, K>

Returns the provided tensor with autodiff enabled.

If the tensor does not yet use autodiff, this associates it with the autodiff backend using the default gradient-checkpointing strategy. If autodiff is already enabled, the tensor and its current strategy are returned unchanged. This operation is idempotent.

This is equivalent to autodiff. from_inner reflects the underlying backend-decorator model, while autodiff describes the operation in terms of the high-level tensor API.

Enabling autodiff does not make a floating-point tensor require gradients. Use require_grad when its gradient should be retained during the backward pass.

pub fn gradient_checkpointing_strategy( &self, ) -> Option<GradientCheckpointingStrategy>

Returns this tensor’s gradient-checkpointing strategy when autodiff is enabled.

pub fn with_gradient_checkpointing_strategy( self, strategy: GradientCheckpointingStrategy, ) -> Tensor<D, K>

Sets the autodiff checkpointing strategy carried by this tensor.

The strategy is normally derived from the device the tensor was created on (see Device::gradient_checkpointing); this method overrides it for a single tensor. Enable autodiff first with autodiff when needed; this method doesn’t enable it automatically.

Panics

Panics if autodiff isn’t enabled. Operations combining tensors that carry different strategies also panic; make sure all operands share the same one.

impl<const D: usize, K> Tensor<D, K>

where K: Basic,

pub fn new(primitive: BridgeTensor) -> Tensor<D, K>

Constructs a new Tensor.

impl<const D: usize, K> Tensor<D, K>

where K: Basic,

pub fn extract(&mut self) -> Tensor<D, K>

Takes ownership of the tensor out of self, leaving an empty zero-shape placeholder tensor in its place.

This is analogous to Option::take / core::mem::take: it lets you obtain an owned Tensor from behind a &mut Tensor so you can call owned operations on it.

pub fn inplace<F>(&mut self, func: F)

where F: FnOnce(Tensor<D, K>) -> Tensor<D, K>,

Executes an operation on the tensor and modifies its value.

Notes

This won’t necessarily reuse the same tensor data/buffer, but it should if there is no other reference pointing to the same tensor.

Wrapping operations with inplace is not an optimization, it’s mainly there if you want to mutate a tensor by using owned operations. A plausible usage would be to update the weights of a mutable model reference.

pub fn rank(&self) -> usize

Returns the number of dimensions of the tensor.

pub fn dtype(&self) -> DType

Returns the tensor primitive data type.

Note

Some element types are encoded in different primitive types depending on the backend (e.g., bool could be encoded as u8 or u32).

pub fn can_mut(&self) -> bool

Whether this tensor’s buffer can be mutated in place — i.e. this handle uniquely owns the allocation, so an in-place op writes it directly instead of copying first (see TensorMetadata::can_mut).

Backends that track buffer ownership (cubecl, fusion, tch) answer precisely from the handle reference count; others conservatively return false — they may alias the buffer, so an in-place write can’t be assumed safe. Useful to assert a hot-path op (e.g. a KV-cache slice_assign) stays in place rather than silently copying.

pub fn empty<S>( shape: S, options: impl Into<TensorCreationOptions>, ) -> Tensor<D, K>

where S: Into<Shape>,

Create an empty tensor of the given shape.

Arguments

• shape: The shape of the tensor.
• device: The device where the tensor will be created.

Example

use burn_tensor::Tensor;

let device = Default::default();
// Create an empty tensor with dimensions [2, 3, 4].
let tensor = Tensor::<3>::empty([2, 3, 4], &device);

pub fn empty_like(&self) -> Tensor<D, K>

Create an empty tensor with the same shape, dtype, and device as the current tensor.

Panics

If the tensor is quantized. This method preserves the input dtype, but quantized tensor creation is not supported.

Example

use burn_tensor::Tensor;

let device = Default::default();
// Create a zeroed tensor with dimensions [2, 3, 4].
let tensor = Tensor::<3>::zeros([2, 3, 4], &device);
// Create an empty tensor with dimensions [2, 3, 4].
let tensor = tensor.empty_like();

pub fn zeros<S>( shape: S, options: impl Into<TensorCreationOptions>, ) -> Tensor<D, K>

where S: Into<Shape>,

Create a tensor of the given shape where each element is zero.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::zeros(Shape::new([2, 3]), &device);
println!("{tensor}");
// [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]

pub fn zeros_like(&self) -> Tensor<D, K>

Returns a new tensor with the same shape, dtype, and device as the current tensor filled with zeros.

Panics

If the tensor is quantized. This method preserves the input dtype, but quantized tensor creation is not supported.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.zeros_like();
println!("{tensor}");
// [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]

pub fn ones<S>( shape: S, options: impl Into<TensorCreationOptions>, ) -> Tensor<D, K>

where S: Into<Shape>,

Create a tensor of the given shape where each element is one.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::ones(Shape::new([2, 3]), &device);
println!("{tensor}");
// [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]

pub fn ones_like(&self) -> Tensor<D, K>

Returns a new tensor with the same shape, dtype, and device as the current tensor filled with ones.

Panics

If the tensor is quantized. This method preserves the input dtype, but quantized tensor creation is not supported.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.ones_like();
println!("{tensor}");
// [[1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]

pub fn full<S, E>( shape: S, fill_value: E, options: impl Into<TensorCreationOptions>, ) -> Tensor<D, K>

where S: Into<Shape>, E: ElementConversion,

Create a tensor of the given shape where each element is equal to the provided value.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::full(Shape::new([2, 3]), 5.0, &device);
println!("{tensor}");
// [[5.0, 5.0, 5.0], [5.0, 5.0, 5.0]]

pub fn full_like<E>(&self, fill_value: E) -> Tensor<D, K>

where E: ElementConversion,

Returns a new tensor with the same shape, dtype, and device as the current tensor, filled with the provided value.

Panics

If the tensor is quantized. This method preserves the input dtype, but quantized tensor creation is not supported.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.full_like(5.0);
println!("{tensor}");
// [[5.0, 5.0, 5.0], [5.0, 5.0, 5.0]]

pub fn dims(&self) -> [usize; D]

Returns the dimensions of the current tensor.

Example

use burn_tensor::Tensor;

let device = Default::default();
let tensor = Tensor::<3>::ones([2, 3, 4], &device);
let dims = tensor.dims(); // [2, 3, 4]
println!("{dims:?}");

pub fn shape(&self) -> Shape

Returns the shape of the current tensor.

Example

use burn_tensor::Tensor;

let device = Default::default();
let tensor = Tensor::<3>::ones([2, 3, 4], &device);
// Shape { dims: [2, 3, 4] }
let shape = tensor.shape();

pub fn reshape<const D2: usize, S>(self, shape: S) -> Tensor<D2, K>

where S: ReshapeArgs<D2>,

Reshape the tensor to have the given shape.

The tensor has the same data and number of elements as the input.

A -1 in the shape is used to infer the remaining dimensions, e.g.: [2, -1] will reshape the tensor with [2, 3, 4] dimensions to [2, 12].

A 0 in the shape instructs to keep the current dimension from the original tensor, e.g.: [2, 0, 4] will reshape the tensor with [2, 3, 4] dimensions to [2, 3, 4]. This is useful when reshaping tensors with unknown dimensions and combining with -1 to infer the remaining dimensions, e.g. [0, -1] will reshape the tensor with [1, 3, 4] dimensions to [1, 12].

Arguments

• shape: The new shape of the tensor.

Panics

• If the tensor contains more than one -1 in the shape.
• If the tensor contains values that are not positive (other than -1).
• If the shape does not match the number of elements of the original shape.

Example

use burn_tensor::Tensor;

let device = Default::default();
// Create a tensor with dimensions [2, 3, 4]
let tensor = Tensor::<3>::ones([2, 3, 4], &device);
// Reshape it to [2, 12], where 12 is inferred from the number of elements.
let reshaped = tensor.reshape([2, -1]);
println!("{reshaped}");

pub fn transpose(self) -> Tensor<D, K>

Transpose the tensor.

For a 2D tensor, this is the standard matrix transpose. For D > 2, the transpose is applied on the last two dimensions. For example, the transpose of a tensor with shape [1, 2, 3, 4] will have shape [1, 2, 4, 3].

See also permute.

Arguments

• tensor - The tensor to transpose.

Returns

The transposed tensor.

Example

use burn_tensor::Tensor;

let device = Default::default();
// Create a 2D tensor of shape [2, 3]
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);

// Transpose the tensor:
// [[1.0, 5.0], [-2.0, 9.0], [3.0, 6.0]]
// The resulting tensor will have dimensions [3, 2].
let transposed = tensor.transpose();
println!("{transposed}");

pub fn t(self) -> Tensor<D, K>

Alias for transpose.

pub fn swap_dims<Dim1, Dim2>(self, dim1: Dim1, dim2: Dim2) -> Tensor<D, K>

where Dim1: AsIndex, Dim2: AsIndex,

Swaps two dimensions of a tensor.

This is a no-op when dim1 == dim2, assuming both are within bounds.

Arguments

• tensor - The tensor to swap the dimensions of.
• dim1 - The first dimension to swap, supports negative indexing.
• dim2 - The second dimension to swap, supports negative indexing.

Returns

The tensor with the dimensions swapped.

Panics

When dimensions are out of bounds.

Example

use burn_tensor::Tensor;

let device = Default::default();
// Create a 2D tensor of shape [2, 3]
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);

// Swap the dimensions 0 and -1 (equivalent to `tensor.transpose()`):
// [[1.0, 5.0], [-2.0, 9.0], [3.0, 6.0]]
// The resulting tensor will have dimensions [3, 2].
let swapped = tensor.swap_dims(0, -1);
println!("{swapped}");

pub fn permute<Dim>(self, axes: [Dim; D]) -> Tensor<D, K>

where Dim: AsIndex,

Permute the dimensions of the tensor.

This is a no-op when the resolved axes match the current order.

Arguments

• axes - The new order of the dimensions. The length of the axes must be equal to the number of dimensions of the tensor. The values must be unique and in the range of the number of dimensions. The values can be negative, in which case they are used as an offset from the end.

Returns

The tensor with the dimensions permuted.

Example

use burn_tensor::Tensor;

let device = Default::default();
// Create a 2D tensor of shape [3, 2]
let tensor = Tensor::<2>::from_data([[1.0, 5.0], [-2.0, 9.0], [3.0, 6.0]], &device);

// Permute the dimensions 1 and 0:
// [[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]]
// The resulting tensor will have dimensions [3, 2].
let permuted = tensor.permute([1, 0]);
println!("{permuted}");

pub fn movedim<S1, S2>(self, src: S1, dst: S2) -> Tensor<D, K>

where S1: MovedimArgs, S2: MovedimArgs,

Moves the dimension(s) of input at the position(s) in source to the position(s) in destination.

Other dimensions of input that are not explicitly moved remain in their original order and appear at the positions not specified in destination.

Arguments

• src - The dimension(s) to move. The values must be unique and in the range of the number of dimensions. The values can be negative, in which case they are used as an offset from the end.

• dst - Destination positions for each of the original dims. These must also be unique. Negative dimensions are counted from the end.

Panics

• If the source and destination dimensions are not of the same length.
• If the source and destination vectors contain duplicate values.
• If the source and destination vectors contain values that are out of bounds.

Returns

The tensor with the dimensions moved.

Example

use burn_tensor::Tensor;

let device = Default::default();
// Create a 3D tensor of shape [3, 2, 1]
let tensor = Tensor::<3>::from_data([[[1.0], [5.0]], [[-2.0], [9.0]], [[3.0], [6.0]]], &device);

// Move the dimensions 0 and 1:
// [[[1.0], [-2.0], [3.0]], [[5.0], [9.0], [6.0]]]
// The resulting tensor will have dimensions [2, 3, 1].
let moved = tensor.movedim(1, 0);
println!("{moved}");

Note

This is a syntactic sugar for permute. It is used widely enough, so we define a separate Op for it

pub fn flip<const N: usize>(self, axes: [impl AsIndex; N]) -> Tensor<D, K>

Reverse the order of elements in the tensor along the given dimensions.

Arguments

• axes - The dimensions to reverse. The values must be unique and in the range of the number of dimensions. The values can be negative, in which case they are used as an offset from the end.

Returns

The tensor with the axes flipped.

Example

use burn_tensor::Tensor;

let device = Default::default();
// Create a 2D tensor with dimensions [4, 3]
let tensor = Tensor::<2>::from_data(
    [
        [3.0, 4.9, 2.0],
        [2.0, 1.9, 3.0],
        [4.0, 5.9, 8.0],
        [1.4, 5.8, 6.0],
    ],
    &device,
);

// Flip the elements in dimensions 0 and 1:
// [[6.0, 5.8, 1.4],
//  [8.0, 5.9, 4.0],
//  [3.0, 1.9, 2.0],
//  [2.0, 4.9, 3.0]]
// The resulting tensor will have dimensions [4, 3].
let flipped = tensor.flip([0, 1]);
println!("{flipped}");

pub fn flatten<const D2: usize>( self, start_dim: impl AsIndex, end_dim: impl AsIndex, ) -> Tensor<D2, K>

Flatten the tensor along a given range of dimensions.

This function collapses the specified range of dimensions into a single dimension, effectively flattening the tensor in that range.

Arguments

• start_dim: The starting dimension of the range to be flattened, supports negative indexing.
• end_dim: The ending dimension of the range to be flattened (inclusive), supports negative indexing.

Type Parameters

• D2: The resulting number of dimensions in the flattened tensor.

Returns

A new Tensor<D2, K> instance with the specified range of dimensions flattened.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
// Create a 3D tensor with dimensions [2, 3, 4]
let tensor = Tensor::<3>::ones(Shape::new([2, 3, 4]), &device);

// Flatten the tensor from dimensions 1 to 2 (inclusive).
// The resulting tensor will have dimensions [2, 12]
let flattened: Tensor<2> = tensor.flatten(1, 2);
println!("{flattened}");

pub fn squeeze<const D2: usize>(self) -> Tensor<D2, K>

Squeeze the tensor along all dimensions, removing dimensions of size one, and effectively reducing the rank of the tensor.

Type Parameters

• D2: The resulting number of dimensions in the squeezed tensor.

Returns

A new Tensor<D2, K> instance with the specified dimension removed.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
// Create a 4D tensor with dimensions [1, 3, 1, 3]
let tensor = Tensor::<4>::from_data(
    [[[[3.0, 4.9, 2.0]], [[2.0, 1.9, 3.0]], [[4.0, 5.9, 8.0]]]],
    &device,
);

// Squeeze the tensor dimensions.
// The resulting tensor will have dimensions [3, 3].
let squeezed = tensor.squeeze::<2>();
println!("{squeezed}");

pub fn squeeze_dim<const D2: usize>(self, dim: impl AsIndex) -> Tensor<D2, K>

Squeeze the tensor along the given dimension, removing the specified dimension of size one, and effectively reducing the rank of the tensor by one.

Arguments

• dim: The dimension to be squeezed. Supports negative indexing.

Type Parameters

• D2: The resulting number of dimensions in the squeezed tensor.

Panics

If the size in the squeezed dimension is not 1.

Returns

A new Tensor<D2, K> instance with the specified dimension removed.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
// Create a 3D tensor with dimensions [3, 1, 3]
let tensor = Tensor::<3>::from_data(
    [[[3.0, 4.9, 2.0]], [[2.0, 1.9, 3.0]], [[4.0, 5.9, 8.0]]],
    &device,
);

// Squeeze the dimension 1.
// The resulting tensor will have dimensions [3, 3].
let squeezed = tensor.squeeze_dim::<2>(1);
println!("{squeezed}");

pub fn squeeze_dims<const D2: usize>( self, dims: &[impl AsIndex], ) -> Tensor<D2, K>

Removes specified dimensions of size 1 from a tensor’s shape. This function takes a tensor and an array of dimensions (dims) to be squeezed. If dims is provided, only the dimensions specified in this array will be removed. Each dimension in dims should correspond to a size of 1 in the tensor; otherwise, the dimension will not be squeezed. If dims is empty, all single-dimensional entries in the tensor will be removed. If entries in dims are negative, then dimensions will be counted from the back.

Arguments

• dims: The dimension(s) to be squeezed.

Type Parameters

• D2: The resulting number of dimensions in the squeezed tensor.

Returns

A new Tensor<D2, K> instance with the specified dimensions removed.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
// Create a 4D tensor with dimensions [2, 1, 4, 1]
let tensor = Tensor::<4>::ones(Shape::new([2, 1, 4, 1]), &device);

// Squeeze the dimensions 1 and 3.
// The resulting tensor will have dimensions [2, 4].
let squeezed: Tensor<2> = tensor.squeeze_dims(&[1, 3]);
println!("{squeezed}");

pub fn unsqueeze<const D2: usize>(self) -> Tensor<D2, K>

Unsqueeze the current tensor. Create new leading dimensions to fit the given size.

Type Parameters

• D2: The resulting number of dimensions in the unsqueezed tensor.

Panics

If the output size D2 is smaller than the current number of dimensions.

Returns

A new Tensor<D2, K> instance with the specified dimensions added.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
// Create a 2D tensor with dimensions [3, 3]
let tensor = Tensor::<2>::ones(Shape::new([3, 3]), &device);
// Unsqueeze the tensor up to 4 dimensions.
// The resulting tensor will have dimensions [1, 1, 3, 3].
let unsqueezed = tensor.unsqueeze::<4>();
println!("{unsqueezed}");

pub fn unsqueeze_dim<const D2: usize>(self, dim: impl AsIndex) -> Tensor<D2, K>

Creates a new tensor with a dimension of size one inserted at the specified position.

Negative dimensions are counted from the end of the valid insertion positions.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
// Create a 2D tensor with dimensions [3, 3]
let tensor = Tensor::<2>::ones(Shape::new([3, 3]), &device);
// Unsqueeze the dimension 1.
// The resulting tensor will have dimensions [3, 1, 3].
let unsqueezed: Tensor<3> = tensor.unsqueeze_dim(1);
println!("{unsqueezed}");

pub fn unsqueeze_dims<const D2: usize>( self, axes: &[impl AsIndex], ) -> Tensor<D2, K>

Creates a new tensor with added dimensions of size one inserted at the specified indices. The indices can be negative, in which case they are counted from the last to the first dimension. the axes can contain duplicates, in which case the number of dimensions inserted at the index is the number of duplicates.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
// Create a 3D tensor with dimensions [3, 4, 5]
let tensor = Tensor::<3>::ones(Shape::new([3, 4, 5]), &device);
// Unsqueeze the leading dimension (0) once and the trailing dimension (-1) twice.
// The resulting tensor will have dimensions [1, 3, 4, 5, 1, 1].
let unsqueezed: Tensor<6> = tensor.unsqueeze_dims(&[0, -1, -1]);
println!("{unsqueezed}");

pub fn roll_dim<Shift, Dim>(self, shift: Shift, dim: Dim) -> Tensor<D, K>

where Shift: AsIndex, Dim: AsIndex,

Roll operation along a specific dimension; wrapping around the elements.

Parameters

• shift: The roll extent; supports negative values and wraps around.
• dim: The dimension to roll; supports negative indexing.

Returns

A new tensor with the specified dimension rolled by the given shift amount.

pub fn roll<Shift, Dim>(self, shifts: &[Shift], dims: &[Dim]) -> Tensor<D, K>

where Shift: AsIndex, Dim: AsIndex,

Roll operation.

Note: unlike pytorch, dims and shifts must have the same length.

A given dim may be rolled multiple times, and the shifts will be applied sequentially.

Parameters

• shifts: A slice of shifts corresponding to each dimension; supports negative values and wraps around.
• dims: A slice of dimensions to roll; supports negative indexing.

Returns

A new tensor with the specified dimensions rolled by the given shifts.

pub fn slice<S>(self, slices: S) -> Tensor<D, K>

where S: SliceArg,

Returns a tensor containing the elements selected from the given slices.

This method provides flexible tensor slicing with support for various range types, negative indices, and stepped slicing. The method accepts both single slices and arrays of slices, with the s! macro providing convenient syntax for complex patterns.

Arguments

• slices - Can be:
  • A single range for 1D slicing (e.g., 0..5, .., 2..)
  • An array of ranges (e.g., [0..2, 1..4])
  • The s! macro output for advanced slicing with steps
  • a &Vec<Slice> or &[Slice]

Behavior

• Supports partial and full slicing in any number of dimensions
• Handles negative indices by wrapping from the end (-1 is the last element)
• Automatically clamps ranges that exceed tensor dimensions
• Supports stepped slicing for selecting every nth element
• Negative steps reverse the selection order

Panics

• If the number of slices exceeds the tensor’s dimensions
• If a range is descending (e.g., 2..1) or empty (e.g., 1..1) without negative step
• If a step is zero

Examples

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

let device = Default::default();

// Single dimension slicing - no brackets needed!
let tensor = Tensor::<1, burn_tensor::Int>::arange(0..10, &device);
let slice = tensor.clone().slice(2..8);  // Simple range
assert_eq!(slice.try_into_vec_as::<i32>().unwrap(), vec![2, 3, 4, 5, 6, 7]);

// Using s! macro for single dimension with step
let slice = tensor.clone().slice(s![0..10;2]);  // Every 2nd element
assert_eq!(slice.try_into_vec_as::<i32>().unwrap(), vec![0, 2, 4, 6, 8]);

// Reverse a dimension with negative step
let slice = tensor.slice(s![..;-1]);  // Reverse entire tensor
assert_eq!(slice.try_into_vec_as::<i32>().unwrap(), vec![9, 8, 7, 6, 5, 4, 3, 2, 1, 0]);

// Multi-dimensional slicing
let tensor = Tensor::<2>::ones(Shape::new([4, 6]), &device);

// Array syntax for simple ranges
let slice = tensor.clone().slice([1..3, 2..5]);
assert_eq!(slice.dims(), [2, 3]);

// Advanced multi-dimensional with s! macro
let slice = tensor.clone().slice(s![0..4;2, ..;-1]);  // Every 2nd row, reverse columns
assert_eq!(slice.dims(), [2, 6]);

// Complex 3D example with mixed slice types
let tensor = Tensor::<3>::ones(Shape::new([4, 6, 8]), &device);
let slice = tensor.slice(s![1..3, ..;2, -3..]);  // Rows 1-2, every 2nd col, last 3 depth
assert_eq!(slice.dims(), [2, 3, 3]);

// Using negative indices
let tensor = Tensor::<2>::ones(Shape::new([4, 6]), &device);
let slice = tensor.slice(s![-2.., ..-1]);  // Last 2 rows, all but last column
assert_eq!(slice.dims(), [2, 5]);

See Also

• s! - The recommended macro for creating complex slice specifications
• slice_assign - Assign values to a slice
• slice_fill - Fill a slice with a constant value
• slice_dim - Slice a single dimension

pub fn slice_assign<S>(self, slices: S, values: Tensor<D, K>) -> Tensor<D, K>

where S: SliceArg,

Assigns values to a slice of the tensor and returns the updated tensor.

This method supports advanced slicing with steps, including negative steps for reverse assignment. Like slice, it accepts both single slices and arrays, with the s! macro providing powerful syntax for complex patterns.

Arguments

• slices - Slice specification (same format as slice method)
• values - Tensor with values to assign (must match slice dimensions)

Panics

• If slices exceed tensor dimensions
• If values dimensions don’t match the selected slice shape
• If a step is zero

Examples

use burn_tensor::{Tensor, s};

let device = Default::default();

// Simple assignment to a sub-region
let mut tensor = Tensor::<2>::zeros([4, 6], &device);
let values = Tensor::<2>::ones([2, 3], &device);
tensor = tensor.slice_assign([1..3, 2..5], values);
// Now tensor[1..3, 2..5] contains ones

// Single dimension assignment with step
let mut tensor = Tensor::<1>::zeros([10], &device);
let values = Tensor::<1>::ones([5], &device);
tensor = tensor.slice_assign(s![0..10;2], values);
// Now every 2nd element is 1: [1, 0, 1, 0, 1, 0, 1, 0, 1, 0]

// Reverse assignment with negative step
let mut tensor = Tensor::<1>::from_data([0.0, 1.0, 2.0, 3.0, 4.0], &device);
let values = Tensor::<1>::from_data([10.0, 11.0, 12.0, 13.0, 14.0], &device);
tensor = tensor.slice_assign(s![..;-1], values);
// Assigns in reverse: [14, 13, 12, 11, 10]

// Complex multi-dimensional assignment
let mut tensor = Tensor::<3>::zeros([4, 6, 8], &device);
let values = Tensor::<3>::ones([2, 3, 3], &device);
tensor = tensor.slice_assign(s![0..4;2, ..;2, -3..], values);
// Assigns to every 2nd row, every 2nd column, last 3 in depth

// Mixed syntax example
let mut tensor = Tensor::<2>::zeros([8, 8], &device);
let pattern = Tensor::<2>::ones([4, 4], &device);
tensor = tensor.slice_assign(s![..;2, ..;2], pattern);
// Creates a checkerboard pattern with ones

See Also

• s! - The recommended macro for creating complex slice specifications
• slice - Extract a slice from a tensor
• slice_fill - Fill a slice with a constant value

pub fn slice_fill<S, E>(self, slices: S, value: E) -> Tensor<D, K>

where E: Element, S: SliceArg,

Fills a slice of the tensor with a constant value and returns the updated tensor.

Like other slice methods, accepts both single slices and arrays. However, this method currently does not support stepped slicing - use slice_assign with a constant tensor for stepped patterns.

Arguments

• slices - Slice specification (same format as slice method, but no steps)
• value - The value to fill the slice with

Panics

• If slices exceed tensor dimensions
• If any slice has a step != 1 (not yet supported)

Examples

use burn_tensor::{Tensor, s};

let device = Default::default();

// Simple fill for a single dimension
let mut tensor = Tensor::<1>::zeros([10], &device);
tensor = tensor.slice_fill(2..5, 1.0);
// Now tensor is [0, 0, 1, 1, 1, 0, 0, 0, 0, 0]

// Multi-dimensional fill
let mut tensor = Tensor::<2>::zeros([4, 6], &device);
tensor = tensor.slice_fill([1..3, 2..5], -1.0);
// Fills the rectangle at rows 1-2, columns 2-4 with -1

// Using negative indices
let mut tensor = Tensor::<1>::zeros([10], &device);
tensor = tensor.slice_fill(-3.., 2.0);
// Fills the last 3 elements with 2.0

// Complex multi-dimensional example
let mut tensor = Tensor::<3>::ones([4, 6, 8], &device);
tensor = tensor.slice_fill(s![1..3, .., -2..], 0.0);
// Sets rows 1-2, all columns, last 2 in depth to 0

// Stepped slicing is supported
let mut tensor = Tensor::<1>::zeros([10], &device);
tensor = tensor.slice_fill(s![0..10;2], 1.0);
// Now every 2nd element is 1: [1, 0, 1, 0, 1, 0, 1, 0, 1, 0]

See Also

• s! - The macro for creating slice specifications with steps
• slice - Extract a slice from a tensor
• slice_assign - Assign tensor values to a slice

pub fn slice_dim<S>(self, dim: impl AsIndex, slice: S) -> Tensor<D, K>

where S: Into<Slice>,

Returns a new tensor with the specified dimension sliced.

Arguments

• dim: The dimension to slice. Supports negative indexing.
• slice: The slice specification for the dimension. Can be a range (e.g., 2..5), slice with step (via s! macro, e.g., s![0..10;2]), or any type that implements Into<Slice>.

Returns

A new tensor with the specified dimension sliced.

Panics

If the slice is out of bounds for the specified dimension.

Examples

    let tensor = Tensor::<3>::zeros([3, 4, 5], &device);

    // Simple range slicing
    let sliced = tensor.clone().slice_dim(1, 1..3);
    assert_eq!(sliced.shape().as_slice(), [3, 2, 5]);

    // Slicing with step - take every 2nd element
    let sliced = tensor.clone().slice_dim(2, s![0..5;2]);
    assert_eq!(sliced.shape().as_slice(), [3, 4, 3]); // Takes indices 0, 2, 4

    // Reverse slicing with negative step
    let sliced = tensor.clone().slice_dim(1, s![..;-1]);
    assert_eq!(sliced.shape().as_slice(), [3, 4, 5]); // Reverses dimension 1

    // Select from index 2 with step 3
    let sliced = tensor.clone().slice_dim(0, s![2..;3]);
    assert_eq!(sliced.shape().as_slice(), [1, 4, 5]); // Takes only index 2

    // Select single index (reduces dimension to size 1)
    let sliced = tensor.slice_dim(0, 1);
    assert_eq!(sliced.shape().as_slice(), [1, 4, 5]);

See Also

• slice - Slice multiple dimensions simultaneously
• s! - The macro for creating complex slice specifications

pub fn select_dim<const D2: usize>( self, dim: impl AsIndex, index: impl AsIndex, ) -> Tensor<D2, K>

Returns a new tensor selecting a dimension index, and then squeezing that dim.

This is defined as equivalent to t.slice_dim(dim, index).squeeze_dim::<D2>(dim)

Arguments

• dim: The dimension to slice. Supports negative indexing.
• index: the dimension index. Supports negative indexing.

Example

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

let device = Default::default();
let tensor = Tensor::<2>::from_data(
    [
        [1.0, 2.0, 3.0],
        [4.0, 5.0, 6.0],
    ],
    &device,
);

let row1 : Tensor<1> = tensor.clone().select_dim(0, 1);
row1
    .to_data()
    .assert_eq(&TensorData::from([4.0, 5.0, 6.0]), false);

let col1 : Tensor<1> = tensor.clone().select_dim(1, 1);
col1
    .to_data()
    .assert_eq(&TensorData::from([2.0, 5.0]), false);

pub fn device(&self) -> Device

Returns the device of the current tensor.

pub fn to_device(self, device: &Device) -> Tensor<D, K>

Moves the tensor to the target device’s compute resource.

This operation preserves the tensor’s autodiff association and gradient-checkpointing strategy; the target device’s autodiff configuration isn’t applied. Tensors newly created on a device inherit its configuration, while tensors moved to that device retain their existing one. Change it explicitly with autodiff, without_autodiff, or Tensor::with_gradient_checkpointing_strategy.

For tracked floating-point tensors, to_device is a recorded operation, including when the target is the current device. Its output is a non-leaf tensor and cannot retain its own gradient: calling require_grad() on it panics. To create a new leaf on the destination, use tensor.to_device(device).detach().require_grad(), which severs the source connection. Supported transfers between compute backends preserve the graph connection; backward transfers gradients to the original source device. Distributed backward currently requires every distributed parameter to use the same backend as the loss. Incompatible graphs panic before synchronization or gradient computation begins.

Transfers between different compute backends currently read values into host memory as TensorData and upload them to the destination backend. This also applies to gradients transferred during backward and can incur synchronization and data-copy overhead.

Panics

Panics when the backend doesn’t support the requested transfer, including transfers to a capture device with autodiff enabled.

pub fn select(self, dim: impl AsIndex, indices: Tensor<1, Int>) -> Tensor<D, K>

Select tensor elements along the given dimension corresponding to the given indices.

Arguments

• dim - The dimension to select from. Supports negative indexing.
• indices - The indices of the elements to select.

Example

use burn_tensor::{Tensor, Int};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [4.0, 5.0, 6.0]], &device);
let indices = Tensor::<1, Int>::from_data([0], &device);
let tensor = tensor.select(0, indices);
println!("{tensor}");
//  [[1.0, -2.0, 3.0]]

pub fn select_assign( self, dim: impl AsIndex, indices: Tensor<1, Int>, values: Tensor<D, K>, update: IndexingUpdateOp, ) -> Tensor<D, K>

Assign the selected elements along the given dimension corresponding to the given indices from the value tensor to the original tensor using the requested update operation.

Note

• IndexingUpdateOp::Add accumulates values at the selected positions (+=). For booleans, Add is logical or.
• IndexingUpdateOp::Assign replaces values at the selected positions (=), when supported by the backend.
• IndexingUpdateOp::Mul multiplies values at the selected positions (*=), when supported by the backend.

When indices contains duplicate entries, behavior varies by operation:

• For Add, accumulation is supported, though results may be non-deterministic on GPU backends.
• For Assign, duplicate indices result in undefined behavior for both the forward result and the backward gradients.
• For Mul, duplicate indices are multiplied in an unspecified order and backward gradients are undefined.
• For Min and Max, duplicate indices are reduced to the minimum/maximum over all contributions, but the backward gradients are undefined.

For deterministic results and correct gradient calculation across all operations, indices should contain unique entries.

Arguments

• dim - The dimension along which to select. Supports negative indexing.
• indices - The indices to select from the tensor.
• values - The values to assign to the selected indices.
• update - The operation used to update the existing values at the indexed positions.

Example

Example using a 3D tensor:

With IndexingUpdateOp::Add:

input[indices[i], j, k] += values[i, j, k]; // dim = 0 input[i, indices[j], k] += values[i, j, k]; // dim = 1 input[i, j, indices[k]] += values[i, j, k]; // dim = 2

With IndexingUpdateOp::Assign, when supported by the backend, the same indexed locations are replaced instead:

input[indices[i], j, k] = values[i, j, k]; // dim = 0 input[i, indices[j], k] = values[i, j, k]; // dim = 1 input[i, j, indices[k]] = values[i, j, k]; // dim = 2

With IndexingUpdateOp::Mul, when supported by the backend, the selected values are multiplied instead:

input[indices[i], j, k] *= values[i, j, k]; // dim = 0 input[i, indices[j], k] *= values[i, j, k]; // dim = 1 input[i, j, indices[k]] *= values[i, j, k]; // dim = 2

Warning

Not all backends have runtime bound checks for the indices, so make sure they are valid. Otherwise, out of bounds indices could lead to unexpected results instead of panicking.

Panics

If the backend doesn’t support the requested update operation.

pub fn mask_where( self, mask: Tensor<D, Bool>, value: Tensor<D, K>, ) -> Tensor<D, K>

Update the given tensor with the value tensor where the mask is true.

This is similar to mask_fill, however the value is a tensor instead of a scalar.

Example

use burn_tensor::{Tensor, Shape, Bool};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let mask = Tensor::<2, Bool>::from_data([[true, false, true], [false, true, false]], &device);
let value = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
let tensor = tensor.mask_where(mask, value);
println!("{tensor}");
// [[2.0, -2.0, 4.0], [5.0, 2.0, 6.0]]

pub fn mask_fill<E>(self, mask: Tensor<D, Bool>, value: E) -> Tensor<D, K>

where E: ElementConversion,

Update the given tensor with the value where the mask is true.

This is similar to mask_where, however the value is a scalar instead of a tensor.

Example

use burn_tensor::{Tensor, Shape, Bool};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let mask = Tensor::<2, Bool>::from_data([[true, false, true], [false, true, false]], &device);
let tensor = tensor.mask_fill(mask, 3.0);
println!("{tensor}");
// [[3.0, -2.0, 3.0], [5.0, 3.0, 6.0]]

pub fn mask_select(self, mask: Tensor<D, Bool>) -> Tensor<1, K>

Selects the elements of the tensor where mask is true, returned as a 1D tensor in the order of the flattened input tensor.

The mask must have the same shape as the tensor. Unlike torch.masked_select, the mask is not broadcast against the tensor.

Notes

The number of selected elements is data-dependent, so this synchronizes with the device, consistent with argwhere and nonzero. On backends without a native implementation, this reads the entire mask back to the host and computes the indices on the CPU; on lazy backends, it also forces the execution of pending operations.

This makes each call a synchronization point between the host and the device: prefer calling it once on final results (e.g. filtering predictions) rather than inside performance-critical loops.

On an autodiff backend, gradients flow back to the selected elements; positions where mask is false receive a zero gradient.

Panics

• If mask does not have the same shape as the tensor.
• If the mask data cannot be read synchronously (e.g. on wasm); use mask_select_async instead.

Example

use burn_tensor::{Tensor, Bool};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device);
let mask = Tensor::<2, Bool>::from_data([[true, false, true], [false, true, false]], &device);
let selected = tensor.mask_select(mask);
println!("{selected}");
// [1.0, 3.0, 5.0]

pub async fn mask_select_async(self, mask: Tensor<D, Bool>) -> Tensor<1, K>

Selects the elements of the tensor where mask is true, returned as a 1D tensor in the order of the flattened input tensor.

Asynchronous version of mask_select, for backends where the mask cannot be read synchronously (e.g. wasm). The synchronization cost remains; only the waiting is non-blocking.

Panics

If mask does not have the same shape as the tensor.

pub fn gather(self, dim: impl AsIndex, indices: Tensor<D, Int>) -> Tensor<D, K>

Gather tensor elements corresponding to the given indices from the specified dim. The dimension supports negative indexing.

Example using a 3D tensor:

output[i, j, k] = input[indices[i, j, k], j, k]; // dim = 0 output[i, j, k] = input[i, indices[i, j, k], k]; // dim = 1 output[i, j, k] = input[i, j, indices[i, j, k]]; // dim = 2

Notes

The index tensor should have the same shape as the original tensor except for the dim specified.

Warning

Not all backends have runtime bound checks for the indices, so make sure the they are valid. Otherwise, out of bounds indices could lead to unexpected results instead of panicking.

pub fn scatter( self, dim: impl AsIndex, indices: Tensor<D, Int>, values: Tensor<D, K>, update: IndexingUpdateOp, ) -> Tensor<D, K>

Assign the gathered elements corresponding to the given indices along the specified dimension from the value tensor to the original tensor using the requested update operation.

Example using a 3D tensor:

With IndexingUpdateOp::Add:

input[indices[i, j, k], j, k] += values[i, j, k]; // dim = 0 input[i, indices[i, j, k], k] += values[i, j, k]; // dim = 1 input[i, j, indices[i, j, k]] += values[i, j, k]; // dim = 2

With IndexingUpdateOp::Assign, when supported by the backend, the same indexed locations are replaced instead:

input[indices[i, j, k], j, k] = values[i, j, k]; // dim = 0 input[i, indices[i, j, k], k] = values[i, j, k]; // dim = 1 input[i, j, indices[i, j, k]] = values[i, j, k]; // dim = 2

With IndexingUpdateOp::Mul, when supported by the backend, the indexed values are multiplied instead:

input[indices[i, j, k], j, k] *= values[i, j, k]; // dim = 0 input[i, indices[i, j, k], k] *= values[i, j, k]; // dim = 1 input[i, j, indices[i, j, k]] *= values[i, j, k]; // dim = 2

Arguments

• dim - The axis along which to scatter elements. Supports negative indexing.
• indices - The indices of the elements to scatter.
• values - The values to scatter into the tensor.
• update - The operation used to update the existing values at the indexed positions.

Notes

The index tensor should have the same shape as the original tensor except for the specified dimension. The value and index tensors should have the same shape.

When indices contains duplicate entries, behavior varies by operation:

• For Add, accumulation is supported, though results may be non-deterministic on GPU backends.
• For Assign, duplicate indices result in undefined behavior for both the forward result and the backward gradients.
• For Mul, duplicate indices are multiplied in an unspecified order and backward gradients are undefined.
• For Min and Max, duplicate indices are reduced to the minimum/maximum over all contributions, but the backward gradients are undefined.

For deterministic results and correct gradient calculation across all operations, indices should contain unique entries.

Other references to the input tensor will not be modified by this operation.

Warning

Not all backends have runtime bound checks for the indices, so make sure the they are valid. Otherwise, out of bounds indices could lead to unexpected results instead of panicking.

Panics

If the backend doesn’t support the requested update operation.

pub fn scatter_nd<const M: usize, const DV: usize>( self, indices: Tensor<M, Int>, values: Tensor<DV, K>, update: IndexingUpdateOp, ) -> Tensor<D, K>

Multi-dimensional scatter: update the tensor at locations given by indices using the specified update operation.

The size of indices’s last axis (call it K) indexes the leading K dims of self; the batch shape indices.shape[0..M-1] is preserved. values has shape indices.shape[0..M-1] ++ self.shape[K..D]. Constraints: K <= D and M >= 1.

Arguments

• indices - The indices of the elements to scatter.
• values - The values to scatter into the tensor.
• update - The operation used to update the existing values at the indexed positions (e.g., add).

Note

When indices contains duplicate entries, behavior varies by operation:

• For Add, accumulation is supported, though results may be non-deterministic on GPU backends.
• For other operations (Assign, Mul, Min, Max), duplicate indices result in undefined behavior for both the forward result and the backward gradients.

For deterministic results and correct gradient calculation across all operations, indices should contain unique entries.

Warning

Not all backends have runtime bound checks for the indices, so make sure they are valid. Otherwise, out of bounds indices could lead to unexpected results instead of panicking.

pub fn gather_nd<const M: usize, const DV: usize>( self, indices: Tensor<M, Int>, ) -> Tensor<DV, K>

Multi-dimensional gather: collect slices from self at multi-index locations specified by indices.

The size of indices’s last axis (call it K) indexes the leading K dims of self; the batch shape indices.shape[0..M-1] is preserved. The output has shape indices.shape[0..M-1] ++ self.shape[K..D]. Constraints: K <= D and M >= 1.

Warning

Not all backends have runtime bound checks for the indices, so make sure they are valid. Otherwise, out of bounds indices could lead to unexpected results instead of panicking.

pub fn into_data(self) -> TensorData

Converts the data of the current tensor.

Note

For better performance, prefer using a Transaction when reading multiple tensors at once. This may improve laziness, especially if executed on a different thread in native environments.

Returns

The tensor data.

Panics

Panics if the backend fails to read the tensor data or the platform doesn’t support synchronous readback.

pub fn try_into_data(self) -> Result<TensorData, ExecutionError>

Converts the data of the current tensor and returns any error that might have occurred since the last time the device was synchronized.

Note

For better performance, prefer using a Transaction when reading multiple tensors at once. This may improve laziness, especially if executed on a different thread in native environments.

Errors

Returns an error if the backend fails to read the tensor data.

Panics

Panics if the platform doesn’t support synchronous readback.

pub fn to_data(&self) -> TensorData

Converts the data of the current tensor.

Note

For better performance, prefer using a Transaction when reading multiple tensors at once. This may improve laziness, especially if executed on a different thread in native environments.

Returns

The tensor data.

Panics

Panics if the backend fails to read the tensor data or the platform doesn’t support synchronous readback.

pub fn try_to_data(&self) -> Result<TensorData, ExecutionError>

Converts the data of the current tensor.

Note

For better performance, prefer using a Transaction when reading multiple tensors at once. This may improve laziness, especially if executed on a different thread in native environments.

Errors

Returns an error if the backend fails to read the tensor data.

Panics

Panics if the platform doesn’t support synchronous readback.

pub fn into_data_async( self, ) -> impl Future<Output = Result<TensorData, ExecutionError>> + Send

Returns the data of the current tensor.

pub fn to_data_async( &self, ) -> impl Future<Output = Result<TensorData, ExecutionError>> + Send

Returns the data of the current tensor.

pub fn to_data_as<E>(&self) -> TensorData

where E: Element,

Copies the tensor data to host memory and converts it to the dtype represented by E.

The conversion is a no-op if the dtype is the same as the current dtype.

See: Tensor::try_to_data_as.

Returns

A TensorData with dtype E::dtype().

Panics

Panics if synchronous readback isn’t supported, tensor execution or storage access fails, or the data can’t be converted to E.

pub fn try_to_data_as<E>(&self) -> Result<TensorData, TensorReadError>

where E: Element,

Copies the tensor data to host memory and converts it to the dtype represented by E.

By contract, this will yield the same result as tensor.try_to_data()?.try_cast_as::<E>().

The conversion is a no-op if the dtype is the same as the current dtype.

Errors

Returns an error if tensor execution or storage access fails, or the data can’t be converted to E.

Panics

Panics if the platform doesn’t support synchronous readback.

pub fn try_to_vec_as<E>(&self) -> Result<Vec<E>, TensorReadError>

where E: Element,

Copies the tensor data to a host Vec<E>, converting the dtype when necessary.

By contract, this will yield the same result as tensor.try_to_data_as::<E>()?.try_to_vec::<E>().

Errors

Returns an error if tensor execution or storage access fails, or the data can’t be converted to E.

Panics

Panics if the platform doesn’t support synchronous readback.

pub fn to_data_dtype(&self, dtype: DType) -> TensorData

Copies the tensor data to host memory and converts it to dtype.

The conversion is a no-op if the dtype is the same as the current dtype.

See: Tensor::try_to_data_dtype.

Returns

A TensorData with the requested dtype.

Panics

Panics if synchronous readback isn’t supported, tensor execution or storage access fails, or the data can’t be converted to dtype.

pub fn try_to_data_dtype( &self, dtype: DType, ) -> Result<TensorData, TensorReadError>

Copies the tensor data to host memory and converts it to dtype.

By contract, this will yield the same result as tensor.try_to_data()?.try_cast(dtype).

The conversion is a no-op if the dtype is the same as the current dtype.

Errors

Returns an error if tensor execution or storage access fails, or the data can’t be converted to dtype.

Panics

Panics if the platform doesn’t support synchronous readback.

pub fn into_data_as<E>(self) -> TensorData

where E: Element,

Reads the tensor data into host memory and converts it to the dtype represented by E.

The conversion is a no-op if the dtype is the same as the current dtype.

See: Tensor::try_into_data_as.

Returns

A TensorData with dtype E::dtype().

Panics

Panics if synchronous readback isn’t supported, tensor execution or storage access fails, or the data can’t be converted to E.

pub fn try_into_data_as<E>(self) -> Result<TensorData, TensorReadError>

where E: Element,

Reads the tensor data into host memory and converts it to the dtype represented by E.

By contract, this will yield the same result as tensor.try_into_data()?.try_cast_as::<E>().

The conversion is a no-op if the dtype is the same as the current dtype.

Errors

Returns an error if tensor execution or storage access fails, or the data can’t be converted to E.

Panics

Panics if the platform doesn’t support synchronous readback.

pub fn try_into_vec_as<E>(self) -> Result<Vec<E>, TensorReadError>

where E: Element,

Reads the tensor data into a host Vec<E>, converting the dtype when necessary.

By contract, this will yield the same result as tensor.try_into_data_as::<E>()?.try_into_vec::<E>().

Errors

Returns an error if tensor execution or storage access fails, or the data can’t be converted to E.

Panics

Panics if the platform doesn’t support synchronous readback.

pub fn into_data_dtype(self, dtype: DType) -> TensorData

Reads the tensor data into host memory and converts it to dtype.

The conversion is a no-op if the dtype is the same as the current dtype.

See: Tensor::try_into_data_dtype.

Returns

A TensorData with the requested dtype.

Panics

Panics if synchronous readback isn’t supported, tensor execution or storage access fails, or the data can’t be converted to dtype.

pub fn try_into_data_dtype( self, dtype: DType, ) -> Result<TensorData, TensorReadError>

Reads the tensor data into host memory and converts it to dtype.

By contract, this will yield the same result as tensor.try_into_data()?.try_cast(dtype).

The conversion is a no-op if the dtype is the same as the current dtype.

Errors

Returns an error if tensor execution or storage access fails, or the data can’t be converted to dtype.

Panics

Panics if the platform doesn’t support synchronous readback.

pub fn from_data<T>( data: T, options: impl Into<TensorCreationOptions>, ) -> Tensor<D, K>

where T: Into<TensorData>,

Create a tensor from the given data on the given device.

pub fn repeat_dim(self, dim: impl AsIndex, times: usize) -> Tensor<D, K>

Repeat the tensor along the given dimension.

The output tensor has the same shape, except along the given dimension.

Arguments

• dim: The dimension to repeat. Supports negative indexing.
• times: The number of times to repeat the tensor along the given dimension in the new tensor.

Returns

A new tensor with the given dimension repeated times times.

Example

use burn_tensor::Tensor;

let device = Default::default();
// Create a 2D tensor with dimensions [3, 2]
let tensor = Tensor::<2>::from_data([[3.0, 4.9], [2.0, 1.9], [4.0, 5.9]], &device);

// Repeat the tensor along the dimension 0 twice.
// [[3.0, 4.9], [2.0, 1.9], [4.0, 5.9], [3.0, 4.9], [2.0, 1.9], [4.0, 5.9]]
// The resulting tensor will have dimensions [6, 2].
let repeated = tensor.repeat_dim(0, 2);
println!("{repeated}");

pub fn repeat(self, sizes: &[usize]) -> Tensor<D, K>

Repeat the tensor along the given dimensions.

Arguments

• sizes: Borrowed slice of the number of times to repeat each dimension.

Returns

A new tensor with the given dimensions repeated times times.

Panics

If sizes contains more elements than the number of dimensions.

Example

use burn_tensor::Tensor;

let device = Default::default();
// Create a 2D tensor with dimensions [3, 2]
let tensor = Tensor::<2>::from_data([[3.0, 4.9], [2.0, 1.9], [4.0, 5.9]], &device);

// Repeat the tensor along the dimension 0 twice and the dimension 0 once.
// [[3.0, 4.9], [2.0, 1.9], [4.0, 5.9], [3.0, 4.9], [2.0, 1.9], [4.0, 5.9]]
// The resulting tensor will have dimensions [6, 2].
let repeated = tensor.repeat(&[2, 1]);

pub fn equal(self, other: Tensor<D, K>) -> Tensor<D, Bool>

Applies element-wise equal comparison.

Returns

A boolean tensor that is true where input is equal to other and false elsewhere.

Panics

If the two tensors don’t have the same shape.

Example

use burn_tensor::Tensor;

let device = Default::default();
let t1 = Tensor::<2>::from_data([[2.0, 4.9], [2.0, 1.9], [4.0, 5.9]], &device);
let t2 = Tensor::<2>::from_data([[3.0, 4.9], [2.0, 1.9], [4.0, 5.9]], &device);
// Compare the elements of the two 2D tensors with dimensions [3, 2].
// [[false, true], [true, true], [true, true]]
let equal = t1.equal(t2);
println!("{equal}");

pub fn not_equal(self, other: Tensor<D, K>) -> Tensor<D, Bool>

Applies element-wise non-equality comparison.

Returns

A boolean tensor that is true where input is not equal to other and false elsewhere.

Panics

If the two tensors don’t have the same shape.

Example

use burn_tensor::Tensor;

let device = Default::default();
let t1 = Tensor::<2>::from_data([[2.0, 4.9], [2.0, 1.9], [4.0, 5.9]], &device);
let t2 = Tensor::<2>::from_data([[3.0, 4.9], [2.0, 1.9], [4.0, 5.9]], &device);
// Compare the elements of the two 2D tensors for inequality.
// [[true, false], [false, false], [false, false]]
let not_equal = t1.not_equal(t2);
println!("{not_equal}");

pub fn equal_scalar<E>(self, other: E) -> Tensor<D, Bool>

where E: Element,

Applies element wise equal comparison and returns a boolean tensor.

Arguments

• other - The scalar to compare.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.equal_scalar(3.0);
println!("{tensor}");
// [[false, false, true], [false, false, false]]

pub fn not_equal_scalar<E>(self, other: E) -> Tensor<D, Bool>

where E: Element,

Applies element wise non-equality comparison and returns a boolean tensor.

Arguments

• other - The scalar to compare.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.not_equal_scalar(3.0);
println!("{tensor}");
// [[true, true, false], [true, true, true]]

pub fn equal_elem<E>(self, other: E) -> Tensor<D, Bool>

where E: Element,

Alias for equal_scalar.

pub fn not_equal_elem<E>(self, other: E) -> Tensor<D, Bool>

where E: Element,

Alias for not_equal_scalar.

pub fn cat(tensors: Vec<Tensor<D, K>>, dim: impl AsIndex) -> Tensor<D, K>

Concatenates all tensors into a new one along the given dimension. The dimension supports negative indexing.

Panics

• If dim is higher than the rank.
• If tensors is an empty vector.
• If all tensors don’t have the same shape (the dimension dim is ignored).

Example

use burn_tensor::Tensor;

let device = Default::default();
let t1 = Tensor::<2>::from_data([[3.0, 4.9, 2.0, 1.0], [2.0, 1.9, 3.0, 1.0]], &device);
let t2 = Tensor::<2>::from_data([[4.0, 5.9, 8.0], [1.4, 5.8, 6.0]], &device);

// Concatenate the two tensors with shapes [2, 4] and [2, 3] along the dimension 1.
// [[3.0, 4.9, 2.0, 1.0, 4.0, 5.9, 8.0], [2.0, 1.9, 3.0, 1.0, 1.4, 5.8, 6.0]]
// The resulting tensor will have shape [2, 7].
let concat = Tensor::cat(vec![t1, t2], 1);
println!("{concat}");

pub fn stack<const D2: usize>( tensors: Vec<Tensor<D, K>>, dim: impl AsIndex, ) -> Tensor<D2, K>

Concatenates all tensors into a new one along a new dimension. The dimension supports negative indexing.

Panics

• If all tensors don’t have the same shape.
• If the given dimension is outside the D + 1 valid insertion positions.

Example

use burn_tensor::Tensor;

let device = Default::default();
let t1 = Tensor::<2>::from_data([[3.0, 4.9, 2.0], [2.0, 1.9, 3.0]], &device);
let t2 = Tensor::<2>::from_data([[4.0, 5.9, 8.0], [1.4, 5.8, 6.0]], &device);
let t3 = Tensor::<2>::from_data([[4.0, 5.9, 8.0], [1.4, 5.8, 6.0]], &device);

// Concatenate the three tensors with shape [2, 3] along a new dimension, 0.
// [[[3.0, 4.9, 2.0], [2.0, 1.9, 3.0]],
//  [[4.0, 5.9, 8.0], [1.4, 5.8, 6.0]],
//  [[4.0, 5.9, 8.0], [1.4, 5.8, 6.0]]]
// The resulting tensor will have shape [3, 2, 3].
let stacked= Tensor::stack::<3>(vec![t1, t2, t3], 0);
println!("{stacked}");

pub fn iter_dim(self, dim: impl AsIndex) -> DimIter<D, K>

Iterate over slices of tensors alongside a given dimension. The dimension supports negative indexing.

Panics

If given dimension is greater than or equal to tensor rank.

Returns

A tensor iterator.

Example

use burn_tensor::Tensor;
 let device = Default::default();
 let tensor = Tensor::<2>::from_data([[3.0, 4.9, 2.0], [2.0, 1.9, 3.0]], &device);
 // Given a 2D tensor with dimensions [2, 3], iterate over slices of tensors along the dimension 0.
 let iter = tensor.iter_dim(0);
 for (i,tensor) in iter.enumerate() {
   println!("Tensor {}: {}", i, tensor);
   // Tensor 0: Tensor { data: [[3.0, 4.9, 2.0]], ... }
   // Tensor 1: Tensor { data: [[2.0, 1.9, 3.0]], ... }
}

pub fn narrow( self, dim: impl AsIndex, start: usize, length: usize, ) -> Tensor<D, K>

Returns a new tensor with the given dimension narrowed to the given range. The dimension supports negative indexing.

Panics

• If the dimension is greater than the number of dimensions of the tensor.
• If the given range exceeds the number of elements on the given dimension.

Returns

A new tensor with the given dimension narrowed to the given range.

Example

use burn_tensor::Tensor;

let device = Default::default();
// Create a 2D tensor with dimensions [4, 3]
let tensor = Tensor::<2>::from_data(
    [
        [3.0, 4.9, 2.0],
        [2.0, 1.9, 3.0],
        [6.0, 1.5, 7.0],
        [3.0, 4.9, 9.0],
    ],
    &device,
);
// Narrow the tensor along the dimension 0, keeping 3 elements starting from index 1.
// [[2.0, 1.9, 3.0], [6.0, 1.5, 7.0], [3.0, 4.9, 9.0]]
// The resulting tensor will have dimensions [3, 3].
let narrowed = tensor.narrow(0, 1, 3);
println!("{narrowed}");

pub fn chunk(self, chunks: usize, dim: impl AsIndex) -> Vec<Tensor<D, K>>

Attempts to split the tensor into a specified number of chunks along a given dimension. The dimension supports negative indexing. May return less chunks than requested if the tensor size is not divisible by the number of chunks.

When the given dimension is evenly divisible by the number of chunks, the chunks will be of equal size. Otherwise all chunks will be of equal size except for the last one.

Panics

If the dimension is greater than the number of dimensions of the tensor.

Returns

A vector of tensors.

Example

use burn_tensor::Tensor;

let device = Default::default();
// Create a 2D tensor with dimensions [4, 3]
let tensor = Tensor::<2>::from_data(
    [
        [3.0, 4.9, 2.0],
        [2.0, 1.9, 3.0],
        [6.0, 1.5, 7.0],
        [3.0, 4.9, 9.0],
    ],
    &device,
);
// Split the tensor along the dimension 1 into 2 chunks.
// The first chuck will have shape [4, 2]:
// [[3.0, 4.9], [2.0, 1.9], [6.0, 1.5], [3.0, 4.9]]
// The second chunk will have shape [4, 1]:
// [[2.0], [3.0], [7.0], [9.0]]
let chunks = tensor.chunk(2, 1);
println!("{chunks:?}");

pub fn split(self, split_size: usize, dim: impl AsIndex) -> Vec<Tensor<D, K>>

Splits the tensor into chunks of a specified size along a given dimension. The dimension supports negative indexing. Each chunk is a view of the original tensor.

If the tensor size along the given dimension is not divisible by split_size, then the last chunk will be smaller.

Panics

If the specified dimension to split along is greater than the number of dimensions of the tensor.

Returns

A vector of tensors.

Example

use burn_tensor::Tensor;

let device = Default::default();
// Create a 1D tensor with 5 elements
let tensor = Tensor::<1>::from_data([0.0, 1.0, 2.0, 3.0, 4.0], &device);
// Split the tensor into chunks of size 2 along dimension 0
let chunks = tensor.split(2, 0);
// The result is a vector of tensors:
// [Tensor([0.0, 1.0]), Tensor([2.0, 3.0]), Tensor([4.0])]
println!("{:?}", chunks);

pub fn split_with_sizes( self, split_sizes: Vec<usize>, dim: impl AsIndex, ) -> Vec<Tensor<D, K>>

Splits the tensor into chunks with the specified sizes along a given dimension. The dimension supports negative indexing. Each chunk is a view of the original tensor.

The sizes of the chunks are specified in the split_sizes vector. The sum of the sizes in split_sizes must equal the size of the tensor along the specified dimension.

Panics

If the specified dimension to split along is greater than the number of dimensions of the tensor or if the sum of dim_sizes does not equal the size of the tensor along dim.

Returns

A vector of tensors.

Example

use burn_tensor::Tensor;

let device = Default::default();
// Create a 1D tensor with 5 elements
let tensor = Tensor::<1>::from_data([0.0, 1.0, 2.0, 3.0, 4.0], &device);
// Split the tensor into chunks with sizes [2, 3] along dimension 0
let chunks = tensor.split_with_sizes(vec![2, 3], 0);
// The result is a vector of tensors:
// [Tensor([0.0, 1.0]), Tensor([2.0, 3.0, 4.0])]
println!("{:?}", chunks);

pub fn any(self) -> Tensor<1, Bool>

Tests if any element in the tensor evaluates to True.

Arguments

• tensor - The tensor to test. All input tensor types (Float, Int, Bool) are supported.

Returns

A boolean tensor Tensor<1, Bool> containing a single element, True if any element in the input tensor evaluates to True, False otherwise.

Example

use burn_tensor::{Tensor, Bool};

let device = Default::default();
let tensor = Tensor::<2, Bool>::from_data([[true,false,true],[false,true,false]], &device);
let tensor_two = Tensor::<2, Bool>::from_data([[false,false,false],[false,false,false]], &device);

// Given a 2D tensor with dimensions [2, 3], test if any element in the tensor evaluates to True.
let any_tensor = tensor.any();
println!("{}", any_tensor);
// Tensor { data: [true], ... }

// Given a 2D tensor with dimensions [2, 3], test if any element in the tensor evaluates to True.
let any_tensor_two = tensor_two.any();
println!("{}", any_tensor_two);
// Tensor { data: [false], ... }

pub fn any_dim(self, dim: impl AsIndex) -> Tensor<D, Bool>

Tests if any element in the tensor evaluates to True along a given dimension dim.

Arguments

• tensor - The tensor to test. All input tensor types (Float, Int, Bool) are supported.
• dim - The axis along which to test. Supports negative indexing.

Returns

A boolean tensor Tensor<D, Bool> with the same shape as input tensor, except in the dim axis where the size is 1. The elem in the dim axis is True if any element along this dim in the input evaluates to True, False otherwise.

Example

use burn_tensor::{Tensor, Bool};

let device = Default::default();
let tensor =
    Tensor::<2, Bool>::from_data([[true, false, false], [false, true, false]], &device);
// Check if any element in the tensor evaluates to True along the dimension 1.
// [[true], [true]],
let any_dim = tensor.clone().any_dim(1);
println!("{any_dim}");

pub fn all(self) -> Tensor<1, Bool>

Tests if all elements in the tensor evaluate to True.

Arguments

• tensor - The tensor to test. All input tensor types (Float, Int, Bool) are supported.

Returns

A boolean tensor Tensor<1, Bool> with a single element, True if all elements in the input tensor evaluate to True, False otherwise.

Example

use burn_tensor::{Tensor, Bool};

let device = Default::default();
let tensor =
    Tensor::<2, Bool>::from_data([[true, false, true], [true, true, true]], &device);
// Check if all elements in the tensor evaluate to True (which is not the case).
// [false]
let all = tensor.all();
println!("{all}");

pub fn all_dim(self, dim: impl AsIndex) -> Tensor<D, Bool>

Tests if all elements in the tensor evaluate to True along a given dimension dim.

Arguments

• tensor - The tensor to test. All input tensor types (Float, Int, Bool) are supported.
• dim - The axis along which to test. Supports negative indexing.

Returns

A boolean tensor Tensor<D, Bool> with the same shape as input tensor, except in the dim axis where the size is 1. The elem in the dim axis is True if all elements along this dim in the input evaluates to True, False otherwise.

Example

use burn_tensor::{Tensor, Bool};

let device = Default::default();
let tensor =
    Tensor::<2, Bool>::from_data([[true, true, false], [true, true, true]], &device);
// Check if all elements in the tensor evaluate to True along the dimension 1.
// [[true, true, false]]
let all_dim = tensor.clone().all_dim(0);
println!("{all_dim}");

pub fn into_scalar<E>(self) -> E

where E: Element,

Convert the tensor into a scalar.

Panics

• If the tensor doesn’t have exactly one element.
• If synchronous readback isn’t supported or the backend fails to read the tensor data.
• If the data can’t be converted to E.

Returns

The scalar value of the tensor.

Example

use burn_tensor::Tensor;

let device = Default::default();
let tensor = Tensor::<2>::from_data([[3.0]], &device);
// Convert the tensor with a single element into a scalar.
let scalar: f32 = tensor.into_scalar();
println!("{scalar}");

pub fn try_into_scalar<E>(self) -> Result<E, TensorReadError>

where E: Element,

Converts the tensor into a scalar and returns any error that occurred since the last time the device was synchronized.

Errors

Returns an error if the tensor doesn’t contain exactly one element, the backend fails to read its data, or the data can’t be converted to E.

Panics

Panics if the platform doesn’t support synchronous readback.

Returns

The scalar value of the tensor.

pub async fn into_scalar_async<E>(self) -> Result<E, ExecutionError>

where E: Element,

Convert the tensor into a scalar asynchronously.

Panics

Panics if the tensor doesn’t contain exactly one element or its data can’t be converted to E.

Errors

Returns an error if the backend fails to read the tensor data.

pub async fn try_into_scalar_async<E>(self) -> Result<E, TensorReadError>

where E: Element,

Try to convert the tensor into a scalar asynchronously.

Errors

Returns an error if the tensor doesn’t contain exactly one element, the backend fails to read its data, or the data can’t be converted to E.

pub fn expand<const D2: usize, S>(self, shape: S) -> Tensor<D2, K>

where S: BroadcastArgs<D, D2>,

Broadcast the tensor to the given shape.

Only singleton dimensions can be expanded to a larger size. Other dimensions must have the same size (which can be inferred with -1).

Arguments

• shape - The shape to broadcast the tensor to. Can contain -1 for dimensions that should be inferred. The number of elements in the shape must be greater or equal as the number of dimensions of the tensor.

Panics

If the tensor cannot be broadcasted to the given shape.

Returns

A new tensor with the given shape.

Example

use burn_tensor::Tensor;

let device = Default::default();
// Create a 2D tensor with dimensions [3, 1]
let tensor = Tensor::<2>::from_data([[1.], [2.], [3.]], &device);
// Expand the tensor to a new shape [3, 4]
// [[1.0, 1.0, 1.0, 1.0], [2.0, 2.0, 2.0, 2.0], [3.0, 3.0, 3.0, 3.0]]
let expanded = tensor.expand([3, 4]);
println!("{}", expanded);

pub fn unfold<const D2: usize, I>( self, dim: I, size: usize, step: usize, ) -> Tensor<D2, K>

where I: AsIndex,

Unfold windows along a dimension.

Returns a view of the tensor with all complete windows of size size in dimension dim; where windows are advanced by step at each index.

The number of windows is 0 when shape[dim] < size, and otherwise (shape[dim] - size) / step + 1.

The new view will have the unfolded dimension replaced by two dimensions; one in the position of the original dimension, with size equal to the number of windows, and one appended to the right-most position, with size equal to size.

Warning

For the ndarray backend; this is not a view but a copy with duplicated data.

Arguments

• dim - the dimension to unfold.
• size - the size of each unfolded window.
• step - the step between each window.

Returns

A tensor view with the shape [pre=..., windows, post=..., size].

impl<const D: usize> Tensor<D, Bool>

pub fn from_bool<A>(data: A, device: &Device) -> Tensor<D, Bool>

where A: Into<TensorData>,

Create a boolean tensor from data on the given device.

Arguments

• data - The tensor data.
• device - The device on which the tensor will be allocated.

Returns

A boolean tensor.

Example

use burn_tensor::{Tensor, Bool};

let device = Default::default();
let tensor = Tensor::<2, Bool>::from_bool([[true, false], [false, true]], &device);
println!("{tensor}");

pub fn int(self) -> Tensor<D, Int>

Convert the bool tensor into an int tensor.

Returns

An integer tensor where true is converted to 1 and false to 0.

Example

use burn_tensor::{Tensor, Bool};

let device = Default::default();
let bool_tensor = Tensor::<1, Bool>::from_bool([true, false, true], &device);
let int_tensor = bool_tensor.int();
println!("{int_tensor}"); // [1, 0, 1]

pub fn float(self) -> Tensor<D>

Convert the bool tensor into a float tensor.

Returns

A float tensor where true is converted to 1.0 and false to 0.0.

Example

use burn_tensor::{Tensor, Bool};

let device = Default::default();
let bool_tensor = Tensor::<1, Bool>::from_bool([true, false, true], &device);
let float_tensor = bool_tensor.float();
println!("{float_tensor}"); // [1.0, 0.0, 1.0]

pub fn cast<T>(self, dtype: T) -> Tensor<D, <T as Cast<D, Bool>>::OutputKind>

where T: Cast<D, Bool>,

Converts a bool tensor to the specified data type.

Supports casting to IntDType (producing an int tensor) or FloatDType (producing a float tensor).

Example

use burn_tensor::{Tensor, Bool, IntDType, FloatDType};

let device = Default::default();
let bool_tensor = Tensor::<1, Bool>::from_bool([true, false, true], &device);

// Cast to int
let int_tensor = bool_tensor.clone().cast(IntDType::I64);

// Cast to float
let float_tensor = bool_tensor.cast(FloatDType::F32);

pub fn bool_not(self) -> Tensor<D, Bool>

Inverses boolean values.

Example

use burn_tensor::{Tensor, Bool};

let device = Default::default();
let tensor = Tensor::<2, Bool>::from_bool([[true, false], [false, true]], &device);
let inverted = tensor.bool_not();
println!("{inverted}"); // [[false, true], [true, false]]

pub fn bool_and(self, rhs: Tensor<D, Bool>) -> Tensor<D, Bool>

Performs logical and (&&) on two boolean tensors.

Arguments

• rhs - The right-hand side tensor for the AND operation.

Returns

A boolean tensor where each element is the result of self[i] && rhs[i].

Example

use burn_tensor::{Tensor, Bool};

let device = Default::default();
let a = Tensor::<2, Bool>::from_bool([[true, true], [false, false]], &device);
let b = Tensor::<2, Bool>::from_bool([[true, false], [true, false]], &device);
let result = a.bool_and(b);
println!("{result}"); // [[true, false], [false, false]]

pub fn bool_or(self, rhs: Tensor<D, Bool>) -> Tensor<D, Bool>

Performs logical or (||) on two boolean tensors.

Arguments

• rhs - The right-hand side tensor for the OR operation.

Returns

A boolean tensor where each element is the result of self[i] || rhs[i].

Example

use burn_tensor::{Tensor, Bool};

let device = Default::default();
let a = Tensor::<2, Bool>::from_bool([[true, true], [false, false]], &device);
let b = Tensor::<2, Bool>::from_bool([[true, false], [true, false]], &device);
let result = a.bool_or(b);
println!("{result}"); // [[true, true], [true, false]]

pub fn bool_xor(self, rhs: Tensor<D, Bool>) -> Tensor<D, Bool>

Performs logical xor (^) on two boolean tensors.

Arguments

• rhs - The right-hand side tensor for the XOR operation.

Returns

A boolean tensor where each element is the result of self[i] ^ rhs[i]. Returns true when exactly one of the operands is true.

Example

use burn_tensor::{Tensor, Bool};

let device = Default::default();
let a = Tensor::<2, Bool>::from_bool([[true, true], [false, false]], &device);
let b = Tensor::<2, Bool>::from_bool([[true, false], [true, false]], &device);
let result = a.bool_xor(b);
println!("{result}"); // [[false, true], [true, false]]

pub fn nonzero(self) -> Vec<Tensor<1, Int>>

Compute the indices of true elements in the tensor (i.e., non-zero for boolean tensors).

Returns

A vector of tensors, one for each dimension of the given tensor, containing the indices of the non-zero elements in that dimension.

Example

use burn_tensor::{Tensor, Bool};

let device = Default::default();
let tensor = Tensor::<2, Bool>::from_bool(
    [[true, false, true], [false, true, false], [false, true, false]],
    &device,
);
let indices = tensor.nonzero();
println!("{}", indices[0]); // [0, 0, 1, 2]
println!("{}", indices[1]); // [0, 2, 1, 1]

pub async fn nonzero_async(self) -> Vec<Tensor<1, Int>>

Compute the indices of true elements in the tensor (i.e., non-zero for boolean tensors).

Returns

A vector of tensors, one for each dimension of the given tensor, containing the indices of the non-zero elements in that dimension.

pub fn argwhere(self) -> Tensor<2, Int>

Compute the indices of the elements that are true, grouped by element.

Returns

A tensor containing the indices of all non-zero elements of the given tensor. Each row in the result contains the indices of a non-zero element.

Example

use burn_tensor::{Tensor, Bool};

let device = Default::default();
let tensor = Tensor::<2, Bool>::from_bool(
    [[true, false, true], [false, true, false], [false, true, false]],
    &device,
);
let indices = tensor.argwhere();
println!("{indices}"); // [[0, 0], [0, 2], [1, 1], [2, 1]]

pub async fn argwhere_async(self) -> Tensor<2, Int>

Compute the indices of the elements that are true, grouped by element.

Returns

A tensor containing the indices of all non-zero elements of the given tensor. Each row in the result contains the indices of a non-zero element.

pub fn triu_mask<S>(shape: S, offset: i64, device: &Device) -> Tensor<D, Bool>

where S: Into<Shape>,

Creates a mask for the upper triangle of a matrix, which can be used to fill the specified area with a value.

This function generates a boolean tensor representing the mask of the upper triangle of a matrix.

Arguments

• shape: The shape of the matrix.
• offset: The offset from the diagonal, where 0 means the diagonal, and positive values shift towards the upper triangle.
• device: The device on which the tensor will be allocated.

Returns

Returns a boolean tensor where false indicates the elements of the matrix that are part of the upper triangle taking into account the specified offset. All other elements are true.

Example

use burn_tensor::{Tensor, Bool};

let mask = Tensor::<2, Bool>::triu_mask([3, 3], 0, &Default::default());
println!("{mask}");
// [[false, false, false],
//  [true, false, false],
//  [true, true, false]]

pub fn tril_mask<S>(shape: S, offset: i64, device: &Device) -> Tensor<D, Bool>

where S: Into<Shape>,

Creates a mask for the lower triangle of a matrix, which can be used to fill the specified area with a value.

This function generates a boolean tensor representing the mask of the lower triangle of a matrix.

Arguments

• shape: The shape of the matrix.
• offset: The offset from the diagonal, where 0 means the diagonal, and negative values shift towards the lower triangle.
• device: The device on which the tensor will be allocated.

Returns

Returns a boolean tensor where false indicates the elements of the matrix that are part of the lower triangle taking into account the specified offset. All other elements are true.

Example

use burn_tensor::{Tensor, Bool};

let mask = Tensor::<2, Bool>::tril_mask([3, 3], 0, &Default::default());
println!("{mask}");
// [[false, true, true],
//  [false, false, true],
//  [false, false, false]]

pub fn diag_mask<S>(shape: S, offset: i64, device: &Device) -> Tensor<D, Bool>

where S: Into<Shape>,

Creates a mask for the diagonal of a matrix, which can be used to fill the specified area with a value.

This function generates a boolean tensor representing the mask of the diagonal of a matrix.

Arguments

• shape: The shape of the matrix.
• offset: The offset from the diagonal, where 0 means the diagonal, and positive values shift towards the upper triangle.
• device: The device on which the tensor will be allocated.

Returns

Returns a boolean tensor where false indicates the elements of the matrix that are part of the diagonal. All other elements are true.

Example

use burn_tensor::{Tensor, Bool};

let mask = Tensor::<2, Bool>::diag_mask([3, 3], 0, &Default::default());
println!("{mask}");
// [[false, true, true],
//  [true, false, true],
//  [true, true, false]]

impl<const D: usize, K> Tensor<D, K>

where K: Numeric,

pub fn einsum( equation: &str, operands: impl IntoIterator<Item = EinsumOperand<K>>, ) -> Tensor<D, K>

Computes Einstein summation with a runtime equation.

Supports explicit ("ij,jk->ik") and implicit ("ij,jk") outputs, any number of operands, repeated-label diagonals, and broadcasting. See einsum! for equation syntax and broadcasting rules. The macro also checks literal equations and statically determined ranks at compile time.

Operands are contracted from left to right using multiplication, reductions, and batched matrix multiplication. No optimized contraction path is searched. Float operations retain automatic differentiation.

Burn represents scalars as rank-one tensors with shape [1]: a scalar result requires D = 1, and an empty input subscript accepts shape [1]. Other inputs must have exactly the rank described by their subscripts.

Example

use burn_tensor::Tensor;
let device = Default::default();
let matrix = Tensor::<2>::from_floats([[1., 2.], [3., 4.]], &device);
let vector = Tensor::<1>::from_floats([10., 20.], &device);
let result = Tensor::<1>::einsum("ij,j->i", [matrix.into(), vector.into()]);
assert_eq!(result.into_data().try_to_vec::<f32>().unwrap(), [50., 110.]);

Panics

Panics for an invalid equation, no operands, incompatible ranks or dimensions, unequal repeated-label dimensions, mismatched devices or dtypes, quantized operands, or an incorrect output rank D.

impl<const D: usize> Tensor<D>

pub fn erf(self) -> Tensor<D>

Applies the error function element wise.

y_i = erf(x_i)

pub fn hypot(self, other: Tensor<D>) -> Tensor<D>

Applies hypotenuse operation element wise.

y_i = sqrt(x_i^2 + y_i^2)

pub fn recip(self) -> Tensor<D>

Applies reciprocal operation (or multiplicative inverse) element wise.

y_i = 1/x_i

pub fn deg2rad(self) -> Tensor<D>

Converts each of the elements of the input tensor from angles in degrees to radians.

Example

let tensor_in_radians = tensor.deg2rad();

pub fn rad2deg(self) -> Tensor<D>

Converts each of the elements of the input tensor from angles in radians to degrees.

Example

let tensor_in_degrees = tensor.rad2deg();

pub fn round(self) -> Tensor<D>

Applies element wise round operation.

This function implements the round half to even strategy, with halfway cases rounded to the nearest even integer value.

pub fn floor(self) -> Tensor<D>

Applies element wise floor operation.

pub fn ceil(self) -> Tensor<D>

Applies element wise ceil operation.

pub fn from_floats<A>(floats: A, device: &Device) -> Tensor<D>

where A: Into<TensorData>,

Create a tensor from floats (f32) on a given device.

Example

use burn_tensor::Tensor;

let device = Default::default();
let _ = Tensor::<1>::from_floats([1.0, 2.0], &device);
let _ = Tensor::<2>::from_floats([[1.0, 2.0], [3.0, 4.0]], &device);

pub fn int(self) -> Tensor<D, Int>

Returns a new tensor with the same shape and device as the current tensor and the data cast to Integer.

Example

use burn_tensor::Tensor;

let device = Default::default();
let float_tensor = Tensor::<1>::from_floats([1.0, 2.0], &device);
let int_tensor = float_tensor.int();

pub fn random_like(&self, distribution: Distribution) -> Tensor<D>

Returns a new tensor with the same shape, dtype, and device as the current tensor filled random values sampled from the given distribution.

Panics

If the tensor is quantized. This method preserves the input dtype, but quantized tensor creation is not supported.

pub fn var<I>(self, dim: I) -> Tensor<D>

where I: AsIndex,

Calculate the variance along the given dimension.

Negative dimensions are supported and count from the end.

pub fn var_bias<I>(self, dim: I) -> Tensor<D>

where I: AsIndex,

Calculate the variance along the given dimension without applying the Bessel’s correction.

Negative dimensions are supported and count from the end.

pub fn var_mean<I>(self, dim: I) -> (Tensor<D>, Tensor<D>)

where I: AsIndex,

Calculate the variance along the given dimension and also returns the mean.

Negative dimensions are supported and count from the end.

pub fn var_mean_bias<I>(self, dim: I) -> (Tensor<D>, Tensor<D>)

where I: AsIndex,

Calculate the variance along the given dimension without applying the Bessel’s correction and also returns the mean.

Negative dimensions are supported and count from the end.

pub fn median<I>(self, dim: I) -> Tensor<D>

where I: AsIndex,

Returns the median value along the specified dimension.

The median is not unique for input tensors with an even number of elements in the reduced dimension. In this case, the lower of the two medians is returned, following PyTorch’s behavior.

Note

The current implementation performs a full sort along the specified dimension, which has O(nlog(n)) complexity. Additionally, most backends currently fall back to CPU for the sort operation, which may result in slower performance compared to native GPU operations.

Arguments

• dim - The dimension along which to compute the median. Negative dimensions are supported and count from the end.

Returns

• A tensor containing the median values along the specified dimension.

Example 1

let device = Default::default();
let tensor = Tensor::<2>::from_data(
    [[1.0, 5.0, 3.0, 2.0], [8.0, 4.0, 6.0, 7.0]],
    &device,
);

// Median along dimension 0:
// sorted columns are [1.0, 8.0], [4.0, 5.0], [3.0, 6.0], [2.0, 7.0]
let median = tensor.median(0);
// Result: [[1.0, 4.0, 3.0, 2.0]]

// Median along dimension 1:
// sorted rows are [1.0, 2.0, 3.0, 5.0] and [4.0, 6.0, 7.0, 8.0]
let median = tensor.median(1);
// Result: [[2.0], [6.0]]

Example 2

The median across all elements can be calculated as follows:

// D is the number of dimensions of the tensor
let flattened_tensor: Tensor<1> = tensor.flatten(0, D - 1);

// Calculate median for dim 0 since the tensor has become 1 dimensional
let median = flattened_tensor.median(0);
// Result: [4.0]

pub fn median_with_indices<I>(self, dim: I) -> (Tensor<D>, Tensor<D, Int>)

where I: AsIndex,

Returns the median value along the specified dimension and its index.

The median is not unique for input tensors with an even number of elements in the reduced dimension. In this case, the lower of the two medians is returned, following PyTorch’s behavior.

Note

The current implementation performs a full sort along the specified dimension, which has O(nlog(n)) complexity. Additionally, most backends currently fall back to CPU for the sort operation, which may result in slower performance compared to native GPU operations.

Arguments

• dim - The dimension along which to compute the median. Negative dimensions are supported and count from the end.

Returns

A tuple containing:

• A tensor with the median values.
• A tensor with the indices of the median values in the original tensor.

Example

let device = Default::default();
let tensor = Tensor::<2>::from_data(
    [[1.0, 5.0, 3.0, 2.0], [8.0, 4.0, 6.0, 7.0]],
    &device,
);

// Median along dimension 1:
// sorted rows are [1.0, 2.0, 3.0, 5.0] and [4.0, 6.0, 7.0, 8.0]
let (values, indices) = tensor.median_with_indices(1);
// values: [[2.0], [6.0]], indices: [[3], [2]] (position in the original tensor)

pub fn cast<T>(self, dtype: T) -> Tensor<D, <T as Cast<D, Float>>::OutputKind>

where T: Cast<D, Float>,

Converts a tensor to the specified data type.

Supports both within-kind casting (e.g., FloatDType::F64) and cross-kind casting (e.g., IntDType::I64 to produce an int tensor).

This is a no-op when casting to the current dtype within the same kind.

Example

use burn_tensor::{Tensor, FloatDType, IntDType};

let device = Default::default();
let float_tensor = Tensor::<1>::from_floats([1.0, 2.5], &device);

// Within-kind cast (float to float)
let f64_tensor = float_tensor.clone().cast(FloatDType::F64);

// Cross-kind cast (float to int)
let int_tensor = float_tensor.cast(IntDType::I64);

pub fn detach(self) -> Tensor<D>

Detach the current tensor from the autodiff graph.

The returned tensor keeps its autodiff association but starts a new graph lineage, so previous operations aren’t considered during backward. A leaf tensor also preserves its current require_grad setting. This function does nothing when autodiff isn’t enabled.

pub fn require_grad(self) -> Tensor<D>

Mark the tensor to keep gradients during the backward pass.

This function does nothing when the tensor is quantized. Enabling gradient retention doesn’t enable autodiff; use autodiff first when needed.

Panics

For non-quantized tensors, panics if autodiff is disabled or the tensor is a non-leaf. Enable autodiff with autodiff before requesting gradients. Use detach first to start a new graph lineage.

pub fn is_require_grad(&self) -> bool

Returns whether this tensor’s gradient is retained after backward.

This is distinct from Tensor::is_autodiff, which reports whether operations use an autodiff context, and Tensor::is_tracked(), which reports graph participation when the autodiff feature is enabled.

pub fn set_require_grad(self, require_grad: bool) -> Tensor<D>

Sets whether this tensor’s gradient is retained after backward.

This function does nothing when the tensor is quantized. Setting this to false on a non-leaf tensor starts a new graph lineage, like detach, while leaving gradient retention disabled.

Panics

Panics when setting this to true without autodiff or on a non-leaf tensor. Enable autodiff with autodiff first. Setting this to false on a tensor without autodiff is harmless. Quantized tensors remain unchanged.

pub fn cov<I>(self, dim: I, correction_factor: usize) -> Tensor<D>

where I: AsIndex,

Calculate covaraince matrix between different entries alongside a given dimension.

Arguments

• dim - The dimension along which to calculate the covariance. Negative dimensions are supported and count from the end.
• correction_factor - Is usually 1 for samples and 0 for population.

pub fn quantize( self, scheme: &QuantScheme, qparams: QParams<Tensor<1>>, ) -> Tensor<D>

Convert the tensor to a lower precision data type based on the quantization scheme.

Arguments

• scheme - The quantization scheme.
• qparams - The pre-computed quantization parameters.

Returns

The quantized tensor.

pub fn quantize_dynamic(self, scheme: &QuantScheme) -> Tensor<D>

Dynamically convert the tensor to a lower precision data type based on the quantization scheme.

Arguments

• scheme - The quantization scheme.

Returns

The quantized tensor.

Notes

This uses min-max calibration.

pub fn dequantize(self) -> Tensor<D>

Convert the tensor back to a higher precision data type.

If the tensor is not quantized, its value is simply returned.

Returns

The dequantized tensor.

pub fn is_close( self, other: Tensor<D>, rtol: Option<f64>, atol: Option<f64>, ) -> Tensor<D, Bool>

Checks element wise if the tensor is close to another tensor.

The tolerance is defined by the following equation:

abs(a - b) <= (atol + rtol * abs(b))

where `a` is the first tensor, `b` is the second tensor, `rtol` is the relative tolerance,
and `atol` is the absolute tolerance.

Arguments

• other - The tensor to compare with.
• rtol - Optional relative tolerance. Default is 1e-5; see DEFAULT_RTOL.
• atol - Optional absolute tolerance. Default is 1e-8; see DEFAULT_ATOL.

Returns

A boolean tensor with the same shape as the input tensors.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor2 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor1.is_close(tensor2, None, None);
println!("{tensor}");
// [[true, true, true], [true, true, true]]

pub fn all_close( self, other: Tensor<D>, rtol: Option<f64>, atol: Option<f64>, ) -> bool

Checks if all elements are close to another tensor.

The tolerance is defined by the following equation:

abs(a - b) <= (atol + rtol * abs(b))

where `a` is the first tensor, `b` is the second tensor, `rtol` is the relative tolerance,
and `atol` is the absolute tolerance.

Arguments

• other - The tensor to compare with.
• rtol - Optional relative tolerance. Default is 1e-5; see DEFAULT_RTOL.
• atol - Optional absolute tolerance. Default is 1e-8; see DEFAULT_ATOL.

Returns

A boolean scalar.

Remarks

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor2 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let result = tensor1.all_close(tensor2, None, None);
println!("{}", result);
// true

pub fn is_nan(self) -> Tensor<D, Bool>

Returns a new tensor with boolean elements indicating whether each element of the input is NaN.

Returns

A boolean tensor where true indicates NaN and false indicates a non-NaN value.

Example

use burn_tensor::{Tensor, Bool, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, f64::NAN, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.is_nan();
println!("{tensor}");
// [[false, true, false], [false, false, false]]

pub fn contains_nan(self) -> Tensor<1, Bool>

Checks if the tensor contains any NaN values.

Returns

A boolean tensor with a single element indicating whether the tensor contains any NaN values.

Example

use burn_tensor::{Tensor, Bool, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [f64::NAN, 9.0, 6.0]], &device);
let tensor = tensor.contains_nan();
println!("{tensor}");
// [true]
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.contains_nan();
println!("{tensor}");
// [false]

pub fn is_inf(self) -> Tensor<D, Bool>

Returns a new tensor with boolean elements indicating whether each element of the input is infinite (either +INF or -INF).

Returns

A boolean tensor where true indicates that the value is infinite

Example

use burn_tensor::{Tensor, Bool, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, f64::INFINITY, 3.0], [f64::NAN, 9.0, 6.0]], &device);
let tensor = tensor.is_finite();
println!("{tensor}");
// [[false, true, false], [false, false, false]]

pub fn is_finite(self) -> Tensor<D, Bool>

Returns a new tensor with boolean elements indicating whether each element of the input is finite

Returns

A boolean tensor where true indicates that the value is finite and false indicates either INF, -INF or NAN

Example

use burn_tensor::{Tensor, Bool, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, f64::INFINITY, 3.0], [f64::NAN, 9.0, 6.0]], &device);
let tensor = tensor.is_finite();
println!("{tensor}");
// [[true, false, true], [false, true, true]]

pub fn grid_sample_2d( self, grid: Tensor<D>, options: impl Into<GridSampleOptions>, ) -> Tensor<D>

Samples tensor as a two-dimensional spatial grid of (possibly multi-channel) values, using the given locations in [-1, 1].

Arguments

• grid - A tensor of locations, with shape (N, H_out, W_out, 2). Values are [-1, 1]. A [x = -1, y = -1] means top-left, and [x = 1, y = 1] means bottom-right
• options - Grid sampling options (mode, padding_mode, align_corners)

Returns

A tensor with shape (N, C, H_out, W_out)

Example

use burn_tensor::ops::{GridSampleOptions, GridSamplePaddingMode, InterpolateMode};

// Default options (bilinear, zeros padding, align_corners=false)
let output = tensor.grid_sample_2d(grid, GridSampleOptions::default());

// Custom options
let options = GridSampleOptions::new(InterpolateMode::Bilinear)
    .with_padding_mode(GridSamplePaddingMode::Border)
    .with_align_corners(true);
let output = tensor.grid_sample_2d(grid, options);

pub fn cross<Dim>(self, other: Tensor<D>, dim: Dim) -> Tensor<D>

where Dim: AsIndex,

Computes the cross product of self and another tensor along a given dimension.

Both self and other must have size 3 along the specified dim, because the cross product is only defined in three-dimensional space.

Arguments

• other - The other tensor to take the cross product with.
• dim - The dimension along which to compute the cross product.

Returns

A tensor containing the cross product of self and other along dim.

pub fn powf(self, other: Tensor<D>) -> Tensor<D>

Applies element wise power operation with a float Tensor

Arguments

• other - The tensor to apply the power operation with.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
let tensor = tensor1.powf(tensor2);
println!("{tensor}");
// [[1.0, 8.0, 81.0], [5.0, 81.0, 216.0]]

pub fn powf_scalar<E>(self, other: E) -> Tensor<D>

where E: ElementConversion,

Applies element wise power operation with a float scalar

Arguments

• other - The scalar to apply the power operation with.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.powf_scalar(2.0);
println!("{tensor}");
// [[1.0, 4.0, 9.0], [25.0, 81.0, 36.0]]

impl<const D: usize> Tensor<D>

pub fn categorical(self, num_samples: usize) -> Tensor<D, Int>

Draws samples from a categorical distribution defined by the last dimension of the input tensor.

The last dimension is treated as a (possibly unnormalized) set of weights defining a categorical distribution over categories. All leading dimensions are treated as batch dimensions. The method returns integer indices of the sampled categories.

Arguments

• num_samples - Number of samples to draw per distribution. Must be >= 1.

Panics

Panics if num_samples is 0.

Note

Distributions with all-zero weights produce undefined (NaN-based) sampling results. Callers should ensure each distribution has at least one positive weight.

Returns

An integer tensor with the same shape as the input, except the last dimension is replaced by num_samples, containing sampled category indices in [0, num_categories).

Example

use burn_tensor::Tensor;

let device = Default::default();
let probs = Tensor::<2>::from_floats(
    [[0.0, 1.0, 0.0], [0.0, 0.0, 1.0]],
    &device,
);
let samples = probs.categorical(4);
// First row always samples index 1, second row always samples index 2
println!("{samples}");

impl<const D: usize> Tensor<D>

pub fn is_distributed(&self) -> bool

Returns true if the tensor is marked as distributed.

pub fn set_distributed(self, param_id: ParamId) -> Tensor<D>

Mark the tensor as distributed.

This function does nothing when autodiff or distributed is not enabled.

impl<const D: usize, K> Tensor<D, K>

where K: FloatMath,

pub fn square(self) -> Tensor<D, K>

Applies element wise square operation.

y_i = x_i * x_i

pub fn exp(self) -> Tensor<D, K>

Applies element wise exponential operation.

y = e^x

pub fn log1p(self) -> Tensor<D, K>

Applies element wise natural logarithm of one plus the input tensor.

y_i = log1p(x_i)

pub fn log(self) -> Tensor<D, K>

Applies element wise natural log operation ln.

y_i = log(x_i)

pub fn sqrt(self) -> Tensor<D, K>

Applies element wise square root operation.

pub fn cos(self) -> Tensor<D, K>

Applies element wise cosine operation.

y_i = cos(x_i)

pub fn sin(self) -> Tensor<D, K>

Applies element wise sine operation.

y_i = sin(x_i)

pub fn tan(self) -> Tensor<D, K>

Applies element wise tangent operation.

y_i = tan(x_i)

pub fn cosh(self) -> Tensor<D, K>

Applies element wise hyperbolic cosine operation.

y_i = cosh(x_i)

Example

use burn_tensor::Tensor;

let device = Default::default();

let tensor = Tensor::<1>::from_data([0.0, -1.0, 2.0], &device);
println!("{}", tensor.cosh()); // [1.0, 1.5430, 3.7621]

pub fn sinh(self) -> Tensor<D, K>

Applies element wise hyperbolic sine operation.

y_i = sinh(x_i)

Example

use burn_tensor::Tensor;

let device = Default::default();

let tensor = Tensor::<1>::from_data([0.0, -1.0, 2.0], &device);
println!("{}", tensor.sinh()); // [0.0, -1.1752, 3.6269]

pub fn tanh(self) -> Tensor<D, K>

Applies element wise hyperbolic tangent operation.

y_i = tanh(x_i)

Example

use burn_tensor::Tensor;

let device = Default::default();

let tensor = Tensor::<1>::from_data([0.0, -1.0, 2.0], &device);
println!("{}", tensor.tanh()); // [0.0, -0.7616, 0.9640]

pub fn acos(self) -> Tensor<D, K>

Applies element wise inverse cosine operation.

y_i = acos(x_i)

Example

use burn_tensor::Tensor;

let device = Default::default();

let tensor = Tensor::<1>::from_data([0.0, -1.0, 1.0], &device);
println!("{}", tensor.acos()); // [1.5708, 3.1416, 0.0]

pub fn acosh(self) -> Tensor<D, K>

Applies element wise inverse hyperbolic cosine operation.

y_i = acosh(x_i)

Example

use burn_tensor::Tensor;

let device = Default::default();

let tensor = Tensor::<1>::from_data([1.0, 2.0, 3.0], &device);
println!("{}", tensor.acosh()); // [0.0000, 1.3170, 1.7627]

pub fn asin(self) -> Tensor<D, K>

Applies element wise inverse sine operation.

y_i = asin(x_i)

Example

use burn_tensor::Tensor;

let device = Default::default();

let tensor = Tensor::<1>::from_data([0.0, -1.0, 1.0], &device);
println!("{}", tensor.asin()); // [ 0.0000, -1.5708,  1.5708]

pub fn asinh(self) -> Tensor<D, K>

Applies element wise inverse hyperbolic sine operation.

y_i = asinh(x_i)

Example

use burn_tensor::Tensor;

let device = Default::default();

let tensor = Tensor::<1>::from_data([0.0, -1.0, 1.0], &device);
println!("{}", tensor.asinh()); // [ 0.0000, -0.8814,  0.8814]

pub fn atan(self) -> Tensor<D, K>

Applies element wise inverse tangent operation.

y_i = atan(x_i)

Example

use burn_tensor::Tensor;

let device = Default::default();

let tensor = Tensor::<1>::from_data([0.0, -1.0, 2.0], &device);
println!("{}", tensor.atan()); // [ 0.0, -0.7854,  1.1071]

pub fn atanh(self) -> Tensor<D, K>

Applies element wise inverse hyperbolic tangent operation.

y_i = atanh(x_i)

Example

use burn_tensor::Tensor;

let device = Default::default();

let tensor = Tensor::<1>::from_data([0.0, -0.5, 0.5], &device);
println!("{}", tensor.atanh()); // [ 0.0, -0.5493,  0.5493]

pub fn atan2(self, other: Tensor<D, K>) -> Tensor<D, K>

Applies element wise inverse tangent operation using the signs of arguments to determine the correct quadrant.

z_i = atan2(y_i, x_i)

Example

use burn_tensor::Tensor;

let device = Default::default();

let lhs = Tensor::<1>::from_data([-2.0, 2.0, -2.0], &device);
let rhs = Tensor::<1>::from_data([1.0, -1.0, -1.0], &device);
println!("{}", lhs.atan2(rhs)); // [-1.1071,  2.0344, -2.0344]

impl<const D: usize> Tensor<D>

pub fn fmod(self, other: Tensor<D>) -> Tensor<D>

Computes the floating-point remainder of dividing self by other.

The result has the same sign as self and magnitude less than other. This is equivalent to the IEEE 754 remainder operation.

Special Cases (IEEE 754 compliant)

• If self is ±∞ and other is not NaN, NaN is returned
• If other is ±0 and self is not NaN, NaN is returned
• If other is ±∞ and self is finite, self is returned
• If either argument is NaN, NaN is returned

Arguments

• other - The divisor tensor. Must have the same shape as self.

Returns

A tensor with the same shape where each element is the floating-point remainder.

Example

use burn_tensor::Tensor;

let device = Default::default();
let dividend = Tensor::<1>::from_data([5.3, -5.3, 5.3, -5.3], &device);
let divisor = Tensor::<1>::from_data([2.0, 2.0, -2.0, -2.0], &device);
let result = dividend.fmod(divisor);

// Result: [1.3, -1.3, 1.3, -1.3]

pub fn fmod_scalar(self, scalar: f32) -> Tensor<D>

Computes the floating-point remainder of dividing self by a scalar.

The result has the same sign as self and magnitude less than the scalar.

Special Cases (IEEE 754 compliant)

• If self is ±∞ and scalar is not NaN, NaN is returned
• If scalar is ±0 and self is not NaN, NaN is returned
• If scalar is ±∞ and self is finite, self is returned
• If either argument is NaN, NaN is returned

Arguments

• scalar - The scalar divisor.

Returns

A tensor with the same shape where each element is the floating-point remainder.

Example

use burn_tensor::Tensor;

let device = Default::default();
let tensor = Tensor::<1>::from_data([5.3, -5.3, 7.5, -7.5], &device);
let result = tensor.fmod_scalar(2.0);

// Result: [1.3, -1.3, 1.5, -1.5]

impl Tensor<1, Int>

pub fn arange( range: Range<i64>, options: impl Into<TensorCreationOptions>, ) -> Tensor<1, Int>

Returns a new integer tensor on the specified device.

Arguments

• range - The range of values to generate.
• device - The device to create the tensor on.

pub fn arange_step( range: Range<i64>, step: usize, options: impl Into<TensorCreationOptions>, ) -> Tensor<1, Int>

Returns a new integer tensor on the specified device.

Arguments

• range - The range of values to generate.
• step - The step between each value.

impl<const D: usize> Tensor<D, Int>

pub fn from_ints<A>(ints: A, device: &Device) -> Tensor<D, Int>

where A: Into<TensorData>,

Create a tensor from integers (i32), placing it on a given device.

Example

use burn_tensor::{Tensor, Int};

let device = Default::default();
let _x: Tensor<1, Int> = Tensor::from_ints([1, 2], &device);
let _y: Tensor<2, Int> = Tensor::from_ints([[1, 2], [3, 4]], &device);

pub fn float(self) -> Tensor<D>

Returns a new tensor with the same shape and device as the current tensor and the data cast to Float.

Example

use burn_tensor::{Int, Tensor};

let device = Default::default();
let int_tensor = Tensor::<1, Int>::arange(0..5, &device);
let float_tensor = int_tensor.float();

pub fn cartesian_grid<S, const D2: usize>( shape: S, device: &Device, ) -> Tensor<D2, Int>

where S: Into<Shape>,

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.

Arguments

• shape - The shape specifying the dimensions of the tensor.
• device - The device to create the tensor on.

Panics

Panics if D2 is not equal to D+1.

Examples

   use burn_tensor::Int;
   use burn_tensor::{Shape, Tensor};
let device = Default::default();
let result: Tensor<3, _> = Tensor::<2, Int>::cartesian_grid([2, 3], &device);
println!("{}", result);

pub fn square(self) -> Tensor<D, Int>

Applies element wise square operation.

y_i = x_i * x_i

pub fn bitwise_and(self, other: Tensor<D, Int>) -> Tensor<D, Int>

Applies the bitwise logical and operation with each bit representing the integer.

pub fn bitwise_or(self, other: Tensor<D, Int>) -> Tensor<D, Int>

Applies the bitwise logical or operation with another tensor.

pub fn bitwise_xor(self, other: Tensor<D, Int>) -> Tensor<D, Int>

Applies the bitwise logical xor operation with another tensor.

pub fn bitwise_not(self) -> Tensor<D, Int>

Applies the bitwise logical not operation.

pub fn bitwise_and_scalar(self, other: impl ElementConversion) -> Tensor<D, Int>

Applies the bitwise logical and operation with each bit in the scalar and the integers in the tensor.

pub fn bitwise_or_scalar(self, other: impl ElementConversion) -> Tensor<D, Int>

Applies the bitwise logical or operation with each bit in the scalar and the integers in the tensor.

pub fn bitwise_xor_scalar(self, other: impl ElementConversion) -> Tensor<D, Int>

Applies bitwise logical xor operation with each bit in the scalar and the integers in the tensor.

pub fn bitwise_left_shift(self, other: Tensor<D, Int>) -> Tensor<D, Int>

Applies the bitwise left shift operation with the integers in the tensor.

pub fn bitwise_right_shift(self, other: Tensor<D, Int>) -> Tensor<D, Int>

Applies the bitwise right shift operation with the integers in the tensor.

pub fn bitwise_left_shift_scalar( self, other: impl ElementConversion, ) -> Tensor<D, Int>

Applies the bitwise left shift operation with the scalar.

pub fn bitwise_right_shift_scalar( self, other: impl ElementConversion, ) -> Tensor<D, Int>

Applies the bitwise right shift operation with the scalar.

pub fn cast<T>(self, dtype: T) -> Tensor<D, <T as Cast<D, Int>>::OutputKind>

where T: Cast<D, Int>,

Converts a tensor to the specified data type.

Supports both within-kind casting (e.g., IntDType::I64) and cross-kind casting (e.g., FloatDType::F32 to produce a float tensor).

This is a no-op when casting to the current dtype within the same kind.

Example

use burn_tensor::{Tensor, Int, IntDType, FloatDType};

let device = Default::default();
let int_tensor = Tensor::<1, Int>::arange(0..5, &device);

// Within-kind cast (int to int)
let i64_tensor = int_tensor.clone().cast(IntDType::I64);

// Cross-kind cast (int to float)
let float_tensor = int_tensor.cast(FloatDType::F32);

impl<const D: usize, K> Tensor<D, K>

where K: Numeric,

pub fn add(self, other: Tensor<D, K>) -> Tensor<D, K>

Applies element wise addition operation.

y = x2 + x1

Arguments

• other - The tensor to add.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
let tensor = tensor1 + tensor2;
println!("{tensor}");
// [[3.0, 1.0, 7.0], [6.0, 11.0, 9.0]]

pub fn add_scalar<E>(self, other: E) -> Tensor<D, K>

where E: ElementConversion,

Applies element wise addition operation with a scalar.

y = x + s

Arguments

• other - The scalar to add, element wise.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let scalar = 2.0;
let tensor = tensor + scalar;
println!("{tensor}");
// [[3.0, 0.0, 5.0], [7.0, 11.0, 8.0]]

pub fn sub(self, other: Tensor<D, K>) -> Tensor<D, K>

Applies element wise subtraction operation.

y = x2 - x1

Arguments

• other - The tensor to subtract.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
let tensor = tensor1 - tensor2;
println!("{tensor}");
// [[-1.0, -5.0, -1.0], [4.0, 7.0, 3.0]]

pub fn sub_scalar<E>(self, other: E) -> Tensor<D, K>

where E: ElementConversion,

Applies element wise subtraction operation with a scalar.

y = x - s

Arguments

• other - The scalar to subtract, element wise.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let scalar = 2.0;
let tensor = tensor - scalar;
println!("{tensor}");
// [[-1.0, -4.0, 1.0], [3.0, 7.0, 4.0]]

pub fn div(self, other: Tensor<D, K>) -> Tensor<D, K>

Applies element wise division operation.

y = x2 / x1

Arguments

• other - The tensor to divide.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
let tensor = tensor1 / tensor2;
println!("{tensor}");
// [[0.5, -0.6666667, 0.75], [5.0, 4.5, 2.0]]

pub fn div_scalar<E>(self, other: E) -> Tensor<D, K>

where E: ElementConversion,

Applies element wise division operation with a scalar.

y = x / s

Arguments

• other - The scalar to divide, element wise.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let scalar = 2.0;
let tensor = tensor / scalar;
println!("{tensor}");
// [[0.5, -1.0, 1.5], [2.5, 4.5, 3.0]]

pub fn remainder(self, other: Tensor<D, K>) -> Tensor<D, K>

Applies the remainder operation element-wise with another tensor.

For integer tensors, this is floor modulo: a nonzero result has the same sign as the divisor and a magnitude less than the divisor’s magnitude.

Integer portability

Portable behavior requires a nonzero divisor and excludes signed MIN % -1, where MIN is the minimum value of the operand dtype. For these excluded inputs, no particular result or error behavior is guaranteed across backends.

pub fn remainder_scalar<E>(self, other: E) -> Tensor<D, K>

where E: ElementConversion,

Applies element wise the remainder operation with a scalar.

y = x % s

For integer tensors, this is floor modulo: a nonzero result has the same sign as the divisor and a magnitude less than the divisor’s magnitude.

Integer portability

Portable behavior requires a nonzero divisor and excludes signed MIN % -1, where MIN is the minimum value of the operand dtype. For these excluded inputs, no particular result or error behavior is guaranteed across backends.

Arguments

• other - The scalar to divide, element wise.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let scalar = 2.0;
let tensor = tensor1 % scalar;
println!("{tensor}");
// [[1.0, 0.0, 1.0], [1.0, 1.0, 0.0]]

pub fn mul(self, other: Tensor<D, K>) -> Tensor<D, K>

Applies element wise multiplication operation.

y = x2 * x1

Arguments

• other - The tensor to multiply.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
let tensor = tensor1 * tensor2;
println!("{tensor}");
// [[2.0, -6.0, 12.0], [5.0, 18.0, 18.0]]

pub fn mul_scalar<E>(self, other: E) -> Tensor<D, K>

where E: ElementConversion,

Applies element wise multiplication operation with a scalar.

y = x * s

Arguments

• other - The scalar to multiply, element wise.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let scalar = 2.0;
let tensor = tensor * scalar;
println!("{tensor}");
// [[2.0, -4.0, 6.0], [10.0, 18.0, 12.0]]

pub fn neg(self) -> Tensor<D, K>

Switch sign of each element in the tensor.

y = -x

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = -tensor;
println!("{tensor}");
// [[-1.0, 2.0, -3.0], [-5.0, -9.0, -6.0]]

pub fn sign(self) -> Tensor<D, K>

Returns the signs of the elements of the input tensor.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.sign();
println!("{tensor}");
// [[1.0, -1.0, 1.0], [1.0, 1.0, 1.0]]

pub fn mean(self) -> Tensor<1, K>

Aggregate all elements in the tensor with the mean operation.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.mean();
println!("{tensor}");
// [3.6666667]

pub fn sum(self) -> Tensor<1, K>

Aggregate all elements in the tensor with the sum operation.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.sum();
println!("{tensor}");
// [22.0]

pub fn mean_dim<I>(self, dim: I) -> Tensor<D, K>

where I: AsIndex,

Aggregate all elements along the given dimension or axis in the tensor with the mean operation.

Arguments

• dim - The dimension or axis along which to aggregate the elements; supports negative indexing.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let mean = tensor.clone().mean_dim(0);
println!("{mean}");
// [[3.0, 3.5, 4.5]]
let mean = tensor.mean_dim(1);
println!("{mean}");
// [[0.6666667], [6.6666665]]

pub fn mean_dims<I>(self, dims: &[I]) -> Tensor<D, K>

where I: AsIndex,

Aggregate all elements along the given axes in the tensor with the mean operation.

Arguments

• dims - the dimensions to aggregate; supports negative indexing.

Returns

The returned tensor will have the same rank, but the aggregated dimensions will have size 1.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[2.0, 4.0], [6.0, -4.0]], &device);
let tensor = tensor.clone().mean_dims(&[0, 1]);
println!("{tensor}");
// [[2.0]]

pub fn sum_dim<I>(self, dim: I) -> Tensor<D, K>

where I: AsIndex,

Aggregate all elements along the given dimension or axis in the tensor with the sum operation.

Arguments

• dim - The dimension or axis along which to aggregate the elements; supports negative indexing.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let sum = tensor.clone().sum_dim(0);
println!("{sum}");
// [[6.0, 7.0, 9.0]]
let sum = tensor.sum_dim(1);
println!("{sum}");
// [[2.0], [20.0]]

pub fn sum_dims<I>(self, dims: &[I]) -> Tensor<D, K>

where I: AsIndex,

Aggregate all elements along the given axes in the tensor with the sum operation.

Arguments

• dims - the dimensions to aggregate; supports negative indexing.

Returns

The returned tensor will have the same rank, but the aggregated dimensions will have size 1.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.clone().sum_dims(&[0, 1]);
println!("{tensor}");
// [[27]]

pub fn sum_dims_squeeze<const D2: usize, I>(self, dims: &[I]) -> Tensor<D2, K>

where I: AsIndex,

Aggregate and squeeze along the given dimensions.

This is equivalent to tensor.sum_dims(dims).squeeze_dims(dims)

Arguments

• dims - the dimensions to aggregate; supports negative indexing.

Returns

The returned tensor will have the same rank, but the aggregated dimensions will have size 1.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<3>::from_data([
    [[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]],
    [[9.0, 2.0, 5.0], [5.0, 7.0, 7.0]],
], &device);
let tensor = tensor.clone().sum_dims_squeeze::<1, _>(&[0, 1]);
println!("{tensor}");
// [20.0, 16.0, 21.0]

pub fn prod(self) -> Tensor<1, K>

Aggregate all elements in the tensor with the product operation.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.prod();
println!("{tensor}");
// [-1620.0]

pub fn prod_dim<I>(self, dim: I) -> Tensor<D, K>

where I: AsIndex,

Aggregate all elements along the given dimension or axis in the tensor with the product operation.

Arguments

• dim - The dimension or axis along which to aggregate the elements, supports negative indexing.

Returns

The returned tensor will have the same rank, but the aggregated dimension will have size 1.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let prod = tensor.clone().prod_dim(0);
println!("{prod}");
// [[5.0, -18.0, 18.0]]
let prod = tensor.prod_dim(1);
println!("{prod}");
// [[-6.0], [270.0]]

pub fn prod_dims<I>(self, dims: &[I]) -> Tensor<D, K>

where I: AsIndex,

Aggregate all elements along the given axes in the tensor with the prod operation.

Arguments

• dims - the dimensions to aggregate, supports negative indexing.

Returns

The returned tensor will have the same rank, but the aggregated dimensions will have size 1.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.clone().sum_dims(&[0, 1]);
println!("{tensor}");
// [[-1620.0]]

pub fn cumsum<I>(self, dim: I) -> Tensor<D, K>

where I: AsIndex,

Computes the cumulative sum of elements along the given dimension or axis.

Arguments

• dim - The dimension or axis along which to compute the cumulative sum. Negative dimensions are supported and count from the end.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device);
let result = tensor.clone().cumsum(0);
println!("{result}");
// [[1.0, 2.0, 3.0], [5.0, 7.0, 9.0]]
let result = tensor.cumsum(1);
println!("{result}");
// [[1.0, 3.0, 6.0], [4.0, 9.0, 15.0]]

pub fn cumprod<I>(self, dim: I) -> Tensor<D, K>

where I: AsIndex,

Computes the cumulative product of elements along the given dimension or axis.

Arguments

• dim - The dimension or axis along which to compute the cumulative product. Negative dimensions are supported and count from the end.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device);
let result = tensor.clone().cumprod(0);
println!("{result}");
// [[1.0, 2.0, 3.0], [4.0, 10.0, 18.0]]
let result = tensor.cumprod(1);
println!("{result}");
// [[1.0, 2.0, 6.0], [4.0, 20.0, 120.0]]

pub fn abs(self) -> Tensor<D, K>

Apply element wise absolute value operation.

Example

use burn_tensor::{Int, Tensor};

let device = Default::default();
let tensor = Tensor::<2, Int>::from_ints([[1, -2, 3], [4, -5, 6], [7, -8, 9]], &device);
let tensor = tensor.abs();
println!("{tensor}");
// [[1, 2, 3], [4, 5, 6], [7, 8, 9]]

Notes

For signed integer dtypes, this operation uses two’s-complement wraparound semantics, similar to x.wrapping_abs(). For example, abs(i64::MIN) == i64::MIN.

pub fn triu(self, diagonal: i64) -> Tensor<D, K>

Returns the upper triangular part of a matrix (2-D tensor) or batch of matrices input, the other elements of the result tensor out are set to 0.

See also triu_mask.

Arguments

• diagonal - The offset from the diagonal, where 0 means the diagonal, and positive values shift towards the upper triangle.

Example

use burn_tensor::{Int, Tensor};

let device = Default::default();
let tensor = Tensor::<2, Int>::from_ints(
    [
      [1, 2, 3],
      [4, 5, 6],
      [7, 8, 9]
    ],
    &device
);
let tensor = tensor.triu(1);
println!("{tensor}");
// [
//   [0, 2, 3],
//   [0, 0, 6],
//   [0, 0, 0]
// ]

pub fn tril(self, diagonal: i64) -> Tensor<D, K>

Returns the lower triangular part of a matrix (2-D tensor) or batch of matrices input, the other elements of the result tensor out are set to 0.

See also tril_mask.

Arguments

• diagonal - The offset from the diagonal, where 0 means the diagonal, and positive values shift towards the upper triangle.

Example

use burn_tensor::{Int, Tensor};

let device = Default::default();
let tensor = Tensor::<2, Int>::from_ints(
    [
      [1, 2, 3],
      [4, 5, 6],
      [7, 8, 9]
    ],
    &device
);

let tensor = tensor.tril(-1);
println!("{tensor}");
// [
//   [0, 0, 0],
//   [4, 0, 0],
//   [7, 8, 0]
// ]

pub fn powi(self, other: Tensor<D, K>) -> Tensor<D, K>

Applies element wise power operation with a integer Tensor

Arguments

• other - The tensor to apply the power operation with.

Example

use burn_tensor::{Tensor, Shape, Int};

let device = Default::default();
let tensor1 = Tensor::<2, Int>::from_ints([[1, -2, 3], [5, 9, 6]], &device);
let tensor2 = Tensor::<2, Int>::from_ints([[2, 3, 4], [1, 2, 3]], &device);
let tensor = tensor1.powi(tensor2);
println!("{tensor}");
// [[1, -8, 81], [5, 81, 216]]

pub fn powi_scalar<E>(self, other: E) -> Tensor<D, K>

where E: ElementConversion,

Applies element wise power operation with a integer scalar

Arguments

• other - The scalar to apply the power operation with.

Example

use burn_tensor::{Tensor, Shape, Int};

let device = Default::default();
let tensor = Tensor::<2, Int>::from_ints([[1, -2, 3], [5, 9, 6]], &device);
let tensor = tensor.powi_scalar(2);
println!("{tensor}");

// [[1, 4, 9], [25, 81, 36]]
let tensor = Tensor::<2>::from_data([[1.5, -2., 3.], [5., 9., 6.]], &device);
let tensor = tensor.powi_scalar(2);
println!("{tensor}");
// [[2.25, 4., 9.], [25., 81., 36.]]

pub fn bool(self) -> Tensor<D, Bool>

Converts the tensor to a boolean tensor by checking if the elements are non-zero.

Returns

A boolean tensor with the same shape as the input tensor.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [0.0, 9.0, 6.0]], &device);
let tensor = tensor.bool();
println!("{tensor}");
// [
//   [true, true, true],
//   [false, true, true]
// ]

pub fn random<S>( shape: S, distribution: Distribution, options: impl Into<TensorCreationOptions>, ) -> Tensor<D, K>

where S: Into<Shape>,

Create a random tensor of the given shape on the given device where each element is sampled from the given distribution.

See also random_like.

Arguments

• shape - The shape of the tensor.
• distribution - The distribution to sample from.
• device - The device to create the tensor on.

Returns

A new tensor with the given shape and elements sampled from the given distribution.

Example

use burn_tensor::{Tensor, Shape, Distribution};

let device = Default::default();
let distribution = Distribution::Uniform(0.0, 1.0); // Any random value between 0.0 and 1.0
let tensor = Tensor::<2>::random(Shape::new([2, 3]), distribution, &device);
println!("{tensor}");
// [
//   [0.08347523, 0.70498955, 0.60332155],
//   [0.08173251, 0.18028641, 0.97942924]
// ]

pub fn matmul(self, other: Tensor<D, K>) -> Tensor<D, K>

Applies the matrix multiplication operation.

C = AB

Shapes of the form [..., B, 1, K] @ [..., 1, K, N] are reinterpreted as [..., 1, B, K] @ [..., 1, K, N], turning a batched vec-mat into a general matmul, which is often faster.

impl<K> Tensor<1, K>

where K: Numeric,

pub fn dot(self, other: Tensor<1, K>) -> Tensor<1, K>

Calculates the dot product with another tensor.

y = x2.dot(x1)

Arguments

• other - The tensor to compute dot product with.

Notes

Both tensors must have the same number of elements.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor1 = Tensor::<1>::from_data([1.0, 2.0], &device);
let tensor2 = Tensor::<1>::from_data([-2.0, 3.0], &device);
let tensor = tensor1.dot(tensor2);
println!("{tensor}");
// [4]

impl<K> Tensor<2, K>

where K: Numeric,

pub fn eye(size: usize, device: &Device) -> Tensor<2, K>

Creates a new 2D tensor with ones on the diagonal and zeros elsewhere.

Arguments

• size - The size of the square matrix.

impl<const D: usize, K> Tensor<D, K>

where K: Ordered,

pub fn sort<I>(self, dim: I) -> Tensor<D, K>

where I: AsIndex,

Sort the elements by value in ascending order along a given dimension.

This sort is unstable (i.e., may reorder equal elements).

Arguments

• dim - The dimension to sort along. Negative dimensions are supported and count from the end.

Returns

A new tensor with the elements sorted in ascending order along the given dimension.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device);
let sorted = tensor.clone().sort(0);
println!("{sorted}");
// [[5.0, -2.0, 3.0], [12.0, 3.0, 6.0]]
let sorted = tensor.sort(1);
println!("{sorted}");
// [[-2.0, 3.0, 12.0], [3.0, 5.0, 6.0]]

pub fn sort_descending<I>(self, dim: I) -> Tensor<D, K>

where I: AsIndex,

Sort the elements by value in descending order along a given dimension.

This sort is unstable (i.e., may reorder equal elements).

Arguments

• dim - The dimension to sort along. Negative dimensions are supported and count from the end.

Returns

A new tensor with the elements sorted in descending order along the given dimension.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device);
let sorted = tensor.clone().sort_descending(0);
println!("{sorted}");
// [[12.0, 3.0, 6.0], [5.0, -2.0, 3.0]]
let sorted = tensor.sort_descending(1);
println!("{sorted}");
// [[12.0, 3.0, -2.0], [6.0, 5.0, 3.0]]

pub fn sort_with_indices<I>(self, dim: I) -> (Tensor<D, K>, Tensor<D, Int>)

where I: AsIndex,

Sort the elements by value in ascending order along a given dimension. Also returns the indices.

This sort is unstable (i.e., may reorder equal elements).

Arguments

• dim - The dimension to sort along. Negative dimensions are supported and count from the end.

Returns

A tuple containing the sorted tensor and the indices tensor.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device);
let (tensor, indices) = tensor.sort_with_indices(0);
println!("{tensor}");
// [[5.0, -2.0, 3.0], [12.0, 3.0, 6.0]]
println!("{}", indices);
// [[1, 0, 0], [0, 1, 1]]

pub fn sort_descending_with_indices<I>( self, dim: I, ) -> (Tensor<D, K>, Tensor<D, Int>)

where I: AsIndex,

Sort the elements by value in descending order along a given dimension. Also returns the indices.

This sort is unstable (i.e., may reorder equal elements).

Arguments

• dim - The dimension to sort along. Negative dimensions are supported and count from the end.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device);
let (tensor, indices) = tensor.sort_descending_with_indices(0);
println!("{tensor}");
// [[12.0, 3.0, 6.0], [5.0, -2.0, 3.0]]
println!("{}", indices);
// [[0, 1, 1], [1, 0, 0]]

pub fn argsort<I>(self, dim: I) -> Tensor<D, Int>

where I: AsIndex,

Returns the indices that sort the elements by value in ascending order along a given dimension.

This sort is unstable (i.e., may reorder equal elements).

Arguments

• dim - The dimension to sort along. Negative dimensions are supported and count from the end.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device);
let tensor = tensor.argsort(0);
println!("{tensor}");
// [[1, 0, 0], [0, 1, 1]]

pub fn argsort_descending<I>(self, dim: I) -> Tensor<D, Int>

where I: AsIndex,

Returns the indices that sort the elements by value in descending order along a given dimension.

This sort is unstable (i.e., may reorder equal elements).

Arguments

• dim - The dimension to sort along. Negative dimensions are supported and count from the end.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device);
let indices = tensor.clone().argsort_descending(0);
println!("{indices}");
// [[0, 1, 1], [1, 0, 0]]
let indices = tensor.argsort_descending(1);
println!("{indices}");
// [[0, 2, 1], [2, 0, 1]]

pub fn topk<I>(self, k: usize, dim: I) -> Tensor<D, K>

where I: AsIndex,

Returns the k largest elements of the given input tensor along a given dimension.

Arguments

• k - The number of elements to return.
• dim - The dimension to sort along. Negative dimensions are supported and count from the end.

Returns

A new tensor with the k largest elements along the given dimension.

With autodiff, gradients are propagated to the input positions selected during the forward pass. All other input positions receive zero gradient. When values are tied, the gradient follows the indices selected by the backend.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device);
let topk = tensor.clone().topk(2, 0);
println!("{topk}");
// [[12.0, 3.0, 6.0], [5.0, -2.0, 3.0]]
let topk = tensor.topk(1, 1);
println!("{topk}");
// [[12.0], [6.0]]

pub fn topk_with_indices<I>( self, k: usize, dim: I, ) -> (Tensor<D, K>, Tensor<D, Int>)

where I: AsIndex,

Returns the k largest elements of the given input tensor along a given dimension. Also returns the indices.

Arguments

• k - The number of elements to return.
• dim - The dimension to sort along. Negative dimensions are supported and count from the end.

With autodiff, gradients from the returned values are propagated to the input positions identified by the returned indices. All other input positions receive zero gradient.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device);
let (tensor, indices) = tensor.topk_with_indices(2, 0);
println!("{tensor}");
// [[12.0, 3.0, 6.0], [5.0, -2.0, 3.0]]
println!("{}", indices);
// [[0, 1, 1], [1, 0, 0]]
let (tensor, indices) = tensor.topk_with_indices(1, 1);
println!("{tensor}");
// [[12.0], [6.0]]
println!("{indices}");
// [[0], [2]]

pub fn one_hot<const D2: usize>(self, num_classes: usize) -> Tensor<D2, K>

Create a one hot tensor.

Example

use burn_tensor::Tensor;

fn example(){
    let device = Default::default();
    let indices: Tensor<1> = Tensor::from_floats([0.0, 1.0, 2.0, 3.0], &device);
    let one_hot: Tensor<2> = indices.one_hot(4);
    println!("{}", one_hot.to_data());
    // [[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 1.0]]
}

pub fn one_hot_fill<const D2: usize>( self, num_classes: usize, on_value: f32, off_value: f32, axis: impl AsIndex, ) -> Tensor<D2, K>

Create a one-hot encoded tensor with configurable num_classes, on_value, off_value, and axis including high-ranked tensors.

Arguments

• num_classes: The number of classes for the one-hot encoding, which defines the size of the one-hot dimension.
• on_value: The value to assign for active positions (corresponding to indices).
• off_value: The value to assign for inactive positions.
• axis: The axis along which the one-hot dimension is added. Negative dimensions are supported and count from the end.

Returns

A tensor with one additional dimension for the one-hot encoding, where active positions are filled with on_value and others with off_value.

Panics

If the tensor is quantized, since lossy quantized values can’t be trusted as class indices.

Example

use burn_tensor::{Tensor, Float};
let device = Default::default();
let indices: Tensor<2, Float> = Tensor::from_floats([[0., 2.], [1., -1.]], &device);
// One-hot encoding
let tensor: Tensor<3, Float> = indices.one_hot_fill(3, 5.0.into(), 0.0.into(), -1);
println!("{tensor}");
// [[[5.0, 0.0, 0.0],
// [0.0, 0.0, 5.0]],
// [[0.0, 5.0, 0.0],
// [0.0, 0.0, 5.0]]]

pub fn greater(self, other: Tensor<D, K>) -> Tensor<D, Bool>

Applies element wise greater comparison and returns a boolean tensor.

Panics

If the two tensors don’t have the same shape.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor2 = Tensor::<2>::from_data([[1.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
let tensor = tensor1.greater(tensor2);
println!("{tensor}");
// [[false, false, false], [true, true, true]]

pub fn greater_equal(self, other: Tensor<D, K>) -> Tensor<D, Bool>

Applies element wise greater-equal comparison and returns a boolean tensor.

Panics

If the two tensors don’t have the same shape.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor2 = Tensor::<2>::from_data([[1.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
let tensor = tensor1.greater_equal(tensor2);
println!("{tensor}");
// [[true, false, false], [true, true, true]]

pub fn lower(self, other: Tensor<D, K>) -> Tensor<D, Bool>

Applies element wise lower comparison and returns a boolean tensor.

Panics

If the two tensors don’t have the same shape.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor2 = Tensor::<2>::from_data([[1.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
let tensor = tensor1.lower(tensor2);
println!("{tensor}");
// [[false, true, true], [false, false, false]]

pub fn lower_equal(self, other: Tensor<D, K>) -> Tensor<D, Bool>

Applies element wise lower-equal comparison and returns a boolean tensor.

Panics

If the two tensors don’t have the same shape.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor2 = Tensor::<2>::from_data([[1.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
let tensor = tensor1.lower_equal(tensor2);
println!("{tensor}");
// [[true, true, true], [false, false, false]]

pub fn greater_scalar<E>(self, other: E) -> Tensor<D, Bool>

where E: ElementConversion,

Applies greater than other comparison and returns a boolean tensor.

Arguments

• other - The scalar to compare.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.greater_scalar(3.0);
println!("{tensor}");
// [[false, false, true], [true, true, true]]

pub fn greater_equal_scalar<E>(self, other: E) -> Tensor<D, Bool>

where E: ElementConversion,

Applies greater-equal than other comparison and returns a boolean tensor.

Arguments

• other - The scalar to compare.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.greater_equal_scalar(3.0);
println!("{tensor}");
// [[false, false, true], [true, true, true]]

pub fn lower_scalar<E>(self, other: E) -> Tensor<D, Bool>

where E: ElementConversion,

Applies lower than other comparison and returns a boolean tensor.

Arguments

• other - The scalar to compare.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.lower_scalar(3.0);
println!("{tensor}");
// [[true, true, false], [false, false, false]]

pub fn lower_equal_scalar<E>(self, other: E) -> Tensor<D, Bool>

where E: ElementConversion,

Applies lower-equal than other comparison and returns a boolean tensor.

Arguments

• other - The scalar to compare.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.lower_equal_scalar(3.0);
println!("{tensor}");
// [[true, true, true], [false, false, false]]

pub fn greater_elem<E>(self, other: E) -> Tensor<D, Bool>

where E: ElementConversion,

Alias for greater_scalar.

pub fn greater_equal_elem<E>(self, other: E) -> Tensor<D, Bool>

where E: ElementConversion,

Alias for greater_equal_scalar.

pub fn lower_elem<E>(self, other: E) -> Tensor<D, Bool>

where E: ElementConversion,

Alias for lower_scalar.

pub fn lower_equal_elem<E>(self, other: E) -> Tensor<D, Bool>

where E: ElementConversion,

Alias for lower_equal_scalar.

pub fn argmax(self, dim: impl AsIndex) -> Tensor<D, Int>

Applies the argmax function along the given dimension and returns an integer tensor.

Arguments

• dim - The dimension along which to find the maximum value. Negative dimensions are supported and count from the end.

NaN behavior

For floating-point tensors, NaNs take precedence over non-NaN values. If a reduced slice contains multiple NaNs, the lowest coordinate along dim is returned. Non-NaN ties also return the lowest coordinate.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<3>::ones(Shape::new([2, 3, 3]), &device);
let tensor = tensor.argmax(1);
println!("{:?}", tensor.shape());
// Shape { dims: [2, 1, 3] }

let tensor = Tensor::<1>::from_data([3.0, f32::NAN, f32::NAN], &device);
let index: i32 = tensor.argmax(0).into_scalar();
assert_eq!(index, 1);

pub fn argtopk(self, k: usize, dim: impl AsIndex) -> Tensor<D, Int>

Applies the argtopk function along the given dimension and returns an integer tensor.

Arguments

• k - The number of indices to return.
• dim - The dimension along which to find the largest values. Negative dimensions are supported and count from the end.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<3>::ones(Shape::new([2, 3, 3]), &device);
let tensor = tensor.argtopk(1, 2);
println!("{:?}", tensor.shape());

pub fn max(self) -> Tensor<1, K>

Find the maximum value.

For floating-point tensors, the result is NaN if any element is NaN.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.max();
println!("{tensor}");
// [9.0]

let tensor = Tensor::<1>::from_data([1.0, f32::NAN, 3.0], &device);
let value: f32 = tensor.max().into_scalar();
assert!(value.is_nan());

pub fn max_dim_with_indices<I>(self, dim: I) -> (Tensor<D, K>, Tensor<D, Int>)

where I: AsIndex,

Find the maximum value along the given dimension.

Also returns the indices.

For floating-point tensors, a NaN in a reduced slice produces a NaN value. The returned index is the lowest coordinate containing NaN. Non-NaN ties also return the lowest coordinate.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let (tensor, index) = tensor.max_dim_with_indices(0);
// [[5.0, 9.0, 6.0]]
println!("{tensor}");
// [[1, 1, 1]]
println!("{index}");

pub fn max_abs(self) -> Tensor<1, K>

Find the maximum absolute value.

For floating-point tensors, the result is NaN if any element is NaN.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -7.0, 3.0], [5.0, -1.0, 6.0]], &device);
let tensor = tensor.max_abs();
println!("{tensor}");
// [7.0]

pub fn max_pair(self, other: Tensor<D, K>) -> Tensor<D, K>

Finds the maximum pair wise values with another tensor.

Arguments

• other - Other tensor to find maximum elements with

Returns

A tensor with the same shape as the input tensors containing the maximum value found in the input tensors.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
let tensor = tensor1.max_pair(tensor2);
println!("{tensor}");
// [[2.0, 3.0, 4.0], [5.0, 9.0, 6.0]]

pub fn max_abs_dim<I>(self, dim: I) -> Tensor<D, K>

where I: AsIndex,

Find the maximum absolute value along the given dimension.

Arguments

• dim - The dimension or axis along which to aggregate the elements, supports negative indexing.

Returns

The returned tensor will have the same rank, but the aggregated dimension will have size 1.

For floating-point tensors, a reduced slice produces NaN if it contains a NaN.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.max_dim(0);
println!("{tensor}");
// [[5.0, 9.0, 6.0]]

pub fn max_abs_dims<I>(self, dims: &[I]) -> Tensor<D, K>

where I: AsIndex,

Find the maximum absolute value along the given dimensions.

Arguments

• dims - The dimensions or axes along which to aggregate the elements, supports negative indexing.

Returns

The returned tensor will have the same rank, but the aggregated dimensions will have size 1.

For floating-point tensors, a reduced region produces NaN if it contains a NaN.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.max_abs_dims(&[0, 1]);
println!("{tensor}");
// [[9.0]]

Empty dimensions

Unlike ordinary reductions such as Tensor::max_dims, this is a composite operation equivalent to self.abs().max_dims(dims). If dims is empty, no reduction is performed, but the elementwise absolute-value operation is still applied.

pub fn argmin(self, dim: impl AsIndex) -> Tensor<D, Int>

Applies the argmin function along the given dimension and returns an integer tensor.

Arguments

• dim - The dimension along which to find the minimum value. Negative dimensions are supported and count from the end.

NaN behavior

For floating-point tensors, NaNs take precedence over non-NaN values. If a reduced slice contains multiple NaNs, the lowest coordinate along dim is returned. Non-NaN ties also return the lowest coordinate.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<3>::ones(Shape::new([2, 3, 3]), &device);
let tensor = tensor.argmin(1);
println!("{:?}", tensor.shape());
// Shape { dims: [2, 1, 3] }

pub fn min(self) -> Tensor<1, K>

Find the minimum value.

For floating-point tensors, the result is NaN if any element is NaN.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.min();
println!("{tensor}");
// [-2.0]

pub fn min_dim<I>(self, dim: I) -> Tensor<D, K>

where I: AsIndex,

Find the minimum value along the given dimension.

Arguments

• dim - The dimension or axis along which to aggregate the elements; supports negative indexing.

Returns

The returned tensor will have the same rank, but the aggregated dimension will have size 1.

For floating-point tensors, a reduced slice produces NaN if it contains a NaN.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.min_dim(0);
println!("{tensor}");
// [[1.0, -2.0, 3.0]]

pub fn min_dims<I>(self, dims: &[I]) -> Tensor<D, K>

where I: AsIndex,

Find the minimum value along the given dimensions.

Arguments

• dims - The dimensions or axes along which to aggregate the elements; supports negative indexing.

Returns

The returned tensor will have the same rank, but the aggregated dimensions will have size 1.

For floating-point tensors, a reduced region produces NaN if it contains a NaN.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.min_dims(&[0, 1]);
println!("{tensor}");
// [[-2.0]]

pub fn min_dim_with_indices<I>(self, dim: I) -> (Tensor<D, K>, Tensor<D, Int>)

where I: AsIndex,

Find the minimum value along the given dimension.

Also returns the indices.

For floating-point tensors, a NaN in a reduced slice produces a NaN value. The returned index is the lowest coordinate containing NaN. Non-NaN ties also return the lowest coordinate.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[7.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let (tensor, index) = tensor.min_dim_with_indices(0);
println!("{tensor}");
// [[5.0, -2.0, 3.0]]
println!("{}", index);
// [[1, 0, 0]]

pub fn min_pair(self, other: Tensor<D, K>) -> Tensor<D, K>

Finds the minimum pair wise values with another tensor.

Arguments

• other - Other tensor to find minimum elements with

Returns

A tensor with the same shape as the input tensors containing the minimum value found between each element of the two source tensors.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor1 = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor2 = Tensor::<2>::from_data([[2.0, 3.0, 4.0], [1.0, 2.0, 3.0]], &device);
let tensor = tensor1.min_pair(tensor2);
println!("{tensor}");
// [[1.0, -2.0, 3.0], [1.0, 2.0, 3.0]]

pub fn clamp<E>(self, min: E, max: E) -> Tensor<D, K>

where E: ElementConversion,

Clamp element wise between the given min and max values.

For floating-point tensors, if the result is zero and the input and bounds include zeros of opposite signs, either sign of zero may be returned.

Arguments

• min - The minimum value.
• max - The maximum value.

Returns

A new tensor with the values clamped between the given min and max values.

Example

use burn_tensor::{Int, Tensor};

let device = Default::default();
let tensor = Tensor::<2, Int>::from_ints(
 [
  [1, 2, 3],
  [4, 5, 6],
  [7, 8, 9]
 ],
 &device);
 let tensor = tensor.clamp(2, 6);
 println!("{tensor}");
 // [[2, 2, 3], [4, 5, 6], [6, 6, 6]]

pub fn clamp_min<E>(self, min: E) -> Tensor<D, K>

where E: ElementConversion,

Clamp element wise under a minimum value.

Arguments

• tensor - The tensor to clamp.
• min - The minimum value.

Returns

A new tensor with the values clamped under the given min value.

Example

use burn_tensor::{Int, Tensor};

let device = Default::default();
let tensor = Tensor::<2, Int>::from_ints(
[[1, 2, 3], [4, 5, 6], [7, 8, 9]],
&device);
let tensor = tensor.clamp_min(4);
println!("{tensor}");
// [[4, 4, 4], [4, 5, 6], [7, 8, 9]]

pub fn clamp_max<E>(self, max: E) -> Tensor<D, K>

where E: ElementConversion,

Clamp element wise over a maximum value.

Arguments

• tensor - The tensor to clamp.
• max - The maximum value.

Returns

A new tensor with the values clamped over the given max value.

Example

use burn_tensor::{Int, Tensor};

let device = Default::default();
let tensor = Tensor::<2, Int>::from_ints(
[[1, 2, 3], [4, 5, 6], [7, 8, 9]],
&device);
let tensor = tensor.clamp_max(5);
println!("{tensor}");
// [[1, 2, 3], [4, 5, 5], [5, 5, 5]]

pub fn cummin<I>(self, dim: I) -> Tensor<D, K>

where I: AsIndex,

Computes the cumulative minimum of elements along the given dimension or axis.

Arguments

• dim - The dimension or axis along which to compute the cumulative minimum. Negative dimensions are supported and count from the end.

For floating-point tensors, once a NaN is encountered in a scan, the output at that position and every later position in the scan is NaN.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[3.0, 5.0, 2.0], [4.0, 1.0, 6.0]], &device);
let result = tensor.clone().cummin(0);
println!("{result}");
// [[3.0, 5.0, 2.0], [3.0, 1.0, 2.0]]
let result = tensor.cummin(1);
println!("{result}");
// [[3.0, 3.0, 2.0], [4.0, 1.0, 1.0]]

pub fn cummax<I>(self, dim: I) -> Tensor<D, K>

where I: AsIndex,

Computes the cumulative maximum of elements along the given dimension or axis.

Arguments

• dim - The dimension or axis along which to compute the cumulative maximum. Negative dimensions are supported and count from the end.

For floating-point tensors, once a NaN is encountered in a scan, the output at that position and every later position in the scan is NaN.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[3.0, 1.0, 2.0], [4.0, 5.0, 2.0]], &device);
let result = tensor.clone().cummax(0);
println!("{result}");
// [[3.0, 1.0, 2.0], [4.0, 5.0, 2.0]]
let result = tensor.cummax(1);
println!("{result}");
// [[3.0, 3.0, 3.0], [4.0, 5.0, 5.0]]

pub fn max_dim<I>(self, dim: I) -> Tensor<D, K>

where I: AsIndex,

Find the maximum value along the given dimension.

Arguments

• dim - The dimension or axis along which to aggregate the elements; supports negative indexing.

Returns

The returned tensor will have the same rank, but the aggregated dimension will have size 1.

For floating-point tensors, a reduced slice produces NaN if it contains a NaN.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.max_dim(0);
println!("{tensor}");
// [[5.0, 9.0, 6.0]]

pub fn max_dims<I>(self, dims: &[I]) -> Tensor<D, K>

where I: AsIndex,

Find the maximum value along the given dimensions.

Arguments

• dims - The dimensions or axis along which to aggregate the elements; supports negative indexing.

Returns

The returned tensor will have the same rank, but the aggregated dimensions will have size 1.

For floating-point tensors, a reduced region produces NaN if it contains a NaN.

Example

use burn_tensor::{Tensor, Shape};

let device = Default::default();
let tensor = Tensor::<2>::from_data([[1.0, -2.0, 3.0], [5.0, 9.0, 6.0]], &device);
let tensor = tensor.max_dims(&[0, 1]);
println!("{tensor}");
// [[9.0]]

impl<const D: usize, K> Tensor<D, K>

where K: Numeric,

pub fn pad( self, padding: impl IntoPadding<D>, mode: impl Into<PadMode>, ) -> Tensor<D, K>

Pads the tensor using the specified padding mode.

Padding is specified as (before, after) pairs. If fewer pairs than tensor dimensions are provided, they apply to the last N dimensions (unspecified leading dimensions are left unpadded).

For backward compatibility, a (left, right, top, bottom) tuple is also accepted, which pads the last two dimensions.

Arguments

• padding - Padding specification. Accepts:
  • [(before, after); N] fixed-size array of pairs (N <= D)
  • &[(before, after)] slice of pairs per dimension
  • Vec<(before, after)> vector of pairs
  • (left, right, top, bottom) tuple for last-2-dim backward compatibility
• mode - The padding mode: Constant(value), Reflect, or Edge.

Returns

A new tensor with the specified padding applied.

Panics

• Panics if more padding pairs are provided than tensor dimensions.
• Reflect mode panics if padding exceeds dimension_size - 1.
• Edge mode panics if padding is applied to a zero-sized dimension.

Example

use burn_tensor::{Tensor, Shape};
use burn_tensor::ops::PadMode;

let device = Default::default();
let tensor = Tensor::<2>::from_data([[12.0, -2.0, 3.0], [5.0, 3.0, 6.0]], &device);

// Constant padding with value 0.0 (backward-compatible tuple)
let padded = tensor.clone().pad((1, 1, 1, 1), PadMode::Constant(0.0));

// Pad arbitrary dimensions with slice of (before, after) pairs
let padded = tensor.clone().pad([(1, 1), (2, 2)], PadMode::Constant(0.0));

// Pad only the last dimension
let padded = tensor.pad([(1, 1)], PadMode::Reflect);

impl<const D: usize, K> Tensor<D, K>

where K: Basic,

pub fn take<const DI: usize, const DO: usize>( self, dim: impl AsIndex, indices: Tensor<DI, Int>, ) -> Tensor<DO, K>

Takes elements from the tensor along the given dimension using indices of any dimensionality.

This behaves like numpy’s take function. When indices is multi-dimensional, the output shape will be: input.shape[:dim] + indices.shape + input.shape[dim+1:]

Arguments

• dim - The dimension along which to select elements. Supports negative indexing.
• indices - The indices of elements to select. Can be any dimensionality. Must be valid indices in the range [0, dim_size).

Example

use burn_tensor::{Tensor, Int};

let device = Default::default();

// Example with 1D indices
let tensor = Tensor::<2>::from_data([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], &device);
let indices = Tensor::<1, Int>::from_data([2, 0, 1], &device);
let result: Tensor<2> = tensor.clone().take::<1, 2>(-1, indices);  // -1 refers to last dimension
println!("{result}");
// [[3.0, 1.0, 2.0], [6.0, 4.0, 5.0]]

// Example with 2D indices - output will have +1 dimension (2D -> 3D)
let indices_2d = Tensor::<2, Int>::from_data([[0, 2], [1, 0]], &device);
let result: Tensor<3> = tensor.take::<2, 3>(1, indices_2d);
println!("{result}");
// [[[1.0, 3.0], [2.0, 1.0]], [[4.0, 6.0], [5.0, 4.0]]]

impl<const D: usize> Tensor<D>

pub fn trunc(self) -> Tensor<D>

Truncates the tensor element-wise, rounding toward zero.

This function returns a new tensor with the same shape as the input tensor, where each element is truncated toward zero. For positive values, this is equivalent to floor, and for negative values, it’s equivalent to ceil.

Special Cases (IEEE 754 compliant)

• trunc(±0) returns ±0 (preserves sign of zero)
• trunc(±∞) returns ±∞
• trunc(NaN) returns NaN

Returns

A tensor with the same shape where each element has been truncated toward zero.

Example

use burn_tensor::Tensor;

let device = Default::default();
let tensor = Tensor::<1>::from_data([2.3, -1.7, 0.5, -0.5, 3.9], &device);
let truncated = tensor.trunc();

// Result: [2.0, -1.0, 0.0, -0.0, 3.0]

impl<const D: usize, K> Tensor<D, K>

where K: Basic,

pub fn into_bridge(self) -> BridgeTensor

Converts the tensor into its bridge-layer representation.

This is primarily intended for backend extensions, allowing custom operations to be inserted between the high-level tensor API and the dispatch layer before deferring to a concrete backend.

pub fn from_bridge(tensor: BridgeTensor) -> Tensor<D, K>

Reconstructs a tensor from its BridgeTensor bridge representation.

This is the inverse of Tensor::into_bridge and is primarily intended for backend extensions, to wrap the output of a custom operation back into the high-level tensor API.

Panics

Panics if the BridgeTensor variant does not match the tensor kind K (e.g. passing an int BridgeTensor when K is Float.

pub fn from_dispatch(tensor: DispatchTensor) -> Tensor<D, K>

Converts from a dispatch tensor into a tensor.

Panics

Panics if the dispatch dtype does not match the tensor kind K.

pub fn into_dispatch(self) -> DispatchTensor

Converts the tensor into a dispatch tensor.

pub fn try_into_primitive<B>( self, ) -> Result<<K as BackendPrimitive<B>>::Primitive, PrimitiveConversionError>

where B: Backend, K: BackendPrimitive<B>, DispatchTensor: DispatchKindConversion<B>,

Safely downcasts a Tensor into its low-level backend primitive.

This is primarily intended for backend extensions where interfacing with the direct backend primitive (e.g., B::FloatTensorPrimitive or AutodiffTensor<B>) is necessary.

Examples

// Extract the underlying CubeCL tensor primitive on the `Wgpu` backend
let cube_tensor = tensor.try_into_primitive::<Wgpu>()?;

// For an autodiff tensor, we can get the wrapped CubeCL tensor primitive on the `Autodiff<Wgpu>` backend
let ad_tensor = tensor.try_into_primitive::<Autodiff<Wgpu>>()?;

Errors

Returns a PrimitiveConversionError if the tensor does not currently live on the requested backend B (including Autodiff<B> mismatch).

CubeCL aliases such as Wgpu and Cuda share a backend type. A successful downcast checks that type, not the execution runtime; inspect the primitive’s device before invoking a runtime-specific kernel. With fusion enabled, the primitive is a fusion handle rather than a raw CubeCL tensor.

pub fn from_primitive<B>( primitive: <K as BackendPrimitive<B>>::Primitive, ) -> Tensor<D, K>

where B: Backend, K: BackendPrimitive<B>, DispatchTensor: DispatchKindConversion<B>,

Reconstructs a high-level Tensor from a concrete backend primitive.

This is the inverse of Tensor::try_into_primitive.

Panics

Panics if the tensor kind K does not match the tensor underlying primitive kind.

Trait Implementations

impl<const D: usize, K> Add for Tensor<D, K>

where K: Numeric,

type Output = Tensor<D, K>

The resulting type after applying the + operator.

fn add(self, rhs: Tensor<D, K>) -> <Tensor<D, K> as Add>::Output

Performs the + operation.

impl<E, const D: usize, K> Add<E> for Tensor<D, K>

where E: ElementConversion, K: Numeric,

type Output = Tensor<D, K>

The resulting type after applying the + operator.

fn add(self, other: E) -> <Tensor<D, K> as Add<E>>::Output

Performs the + operation.

impl<const D: usize, K> Add<Tensor<D, K>> for f64

where K: Numeric,

type Output = Tensor<D, K>

The resulting type after applying the + operator.

fn add(self, tensor: Tensor<D, K>) -> <f64 as Add<Tensor<D, K>>>::Output

Performs the + operation.

impl<const D: usize, K> Add<Tensor<D, K>> for u32

where K: Numeric,

type Output = Tensor<D, K>

The resulting type after applying the + operator.

fn add(self, tensor: Tensor<D, K>) -> <u32 as Add<Tensor<D, K>>>::Output

Performs the + operation.

impl<const D: usize> BitAnd for Tensor<D, Bool>

type Output = Tensor<D, Bool>

The resulting type after applying the & operator.

fn bitand(self, tensor: Tensor<D, Bool>) -> <Tensor<D, Bool> as BitAnd>::Output

Performs the & operation.

impl<const D: usize> BitOr for Tensor<D, Bool>

type Output = Tensor<D, Bool>

The resulting type after applying the | operator.

fn bitor(self, tensor: Tensor<D, Bool>) -> <Tensor<D, Bool> as BitOr>::Output

Performs the | operation.

impl<const D: usize> BitXor for Tensor<D, Bool>

type Output = Tensor<D, Bool>

The resulting type after applying the ^ operator.

fn bitxor(self, tensor: Tensor<D, Bool>) -> <Tensor<D, Bool> as BitXor>::Output

Performs the ^ operation.

impl<const D: usize, K> Clone for Tensor<D, K>

where K: Clone + Basic,

fn clone(&self) -> Tensor<D, K>

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl ColorConversion for Tensor<4>

Formulas derived from the OpenCV docs.

fn rgb2gray(self) -> Tensor<4>

Converts a batch of images from the RGB color space to grayscale.

fn gray2rgb(self) -> Tensor<4>

Converts a batch of grayscale images to the RGB color space.

fn rgb2hsv(self) -> Tensor<4>

Converts a batch of images from the RGB color space to the HSV color space.

fn hsv2rgb(self) -> Tensor<4>

Converts a batch of images from the HSV color space to the RGB color space.

impl ConnectedComponents for Tensor<2, Bool>

fn connected_components(self, connectivity: Connectivity) -> Tensor<2, Int>

Computes the connected components labeled image of boolean image with 4 or 8 way connectivity - returns a tensor of the component label of each pixel.

fn connected_components_with_stats( self, connectivity: Connectivity, options: ConnectedStatsOptions, ) -> (Tensor<2, Int>, ConnectedStats)

Computes the connected components labeled image of boolean image with 4 or 8 way connectivity and collects statistics on each component - returns a tensor of the component label of each pixel, along with stats collected for each component.

impl<const D: usize, K> Debug for Tensor<D, K>

where K: Debug + Basic,

fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error>

Formats the value using the given formatter.

impl<'de, const D: usize, K> Deserialize<'de> for Tensor<D, K>

where K: Basic,

fn deserialize<De>( deserializer: De, ) -> Result<Tensor<D, K>, <De as Deserializer<'de>>::Error>

where De: Deserializer<'de>,

Deserialize this value from the given Serde deserializer.

impl<const D: usize, K> Display for Tensor<D, K>

where K: Basic,

Pretty print tensors

fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), Error>

Formats the value using the given formatter.

impl<const D: usize, K> Div for Tensor<D, K>

where K: Numeric,

type Output = Tensor<D, K>

The resulting type after applying the / operator.

fn div(self, rhs: Tensor<D, K>) -> <Tensor<D, K> as Div>::Output

Performs the / operation.

impl<E, const D: usize, K> Div<E> for Tensor<D, K>

where E: ElementConversion, K: Numeric,

type Output = Tensor<D, K>

The resulting type after applying the / operator.

fn div(self, other: E) -> <Tensor<D, K> as Div<E>>::Output

Performs the / operation.

impl<const D: usize> Div<Tensor<D>> for f64

type Output = Tensor<D>

The resulting type after applying the / operator.

fn div(self, tensor: Tensor<D>) -> <f64 as Div<Tensor<D>>>::Output

Performs the / operation.

impl<const D: usize, K> From<&Tensor<D, K>> for EinsumOperand<K>

where K: Numeric,

fn from(tensor: &Tensor<D, K>) -> EinsumOperand<K>

Converts to this type from the input type.

impl<const D: usize, K, T> From<T> for Tensor<D, K>

where K: Basic, T: Into<TensorData>,

fn from(value: T) -> Tensor<D, K>

Converts to this type from the input type.

impl<const D: usize, K> From<Tensor<D, K>> for EinsumOperand<K>

where K: Numeric,

fn from(tensor: Tensor<D, K>) -> EinsumOperand<K>

Converts to this type from the input type.

impl<const D: usize, K> Module for Tensor<D, K>

where K: Autodiff,

fn visit<V>(&self, _visitor: &mut V)

where V: ModuleVisitor,

Visit each parameter and module-owned control value with a visitor.

fn map<M>(self, _mapper: &mut M) -> Tensor<D, K>

where M: ModuleMapper,

Map each parameter and module-owned control value with a mapper.

fn to_device(self, device: &Device) -> Tensor<D, K>

Move the module and all of its sub-modules to the given device.

fn fork(self, device: &Device) -> Tensor<D, K>

Fork the module and all of its sub-modules to the given device.

fn collect_devices(&self, devices: Vec<Device>) -> Vec<Device>

Return all the devices found in the underneath module tree added to the given vector without duplicates.

fn materialize(self) -> Tensor<D, K>

Fold parameter reparameterizations into their effective values and remove them.

fn valid(&self) -> Tensor<D, K>

Returns a validation snapshot without autodiff or active training flags.

fn train(self) -> Tensor<D, K>

Enables autodiff and restores configured training state throughout the module.

fn devices(&self) -> Vec<Device>

Return all the devices found in the underneath module tree without duplicates.

fn set_require_grad(self, require_grad: bool) -> Self

Set whether every floating-point tensor parameter in the module tree requires gradients.

fn set_require_grad_group(self, group: ParamGroup, require_grad: bool) -> Self

Set whether matched floating-point tensor parameters require gradients.

fn no_grad(self) -> Self

Disable gradients for every floating-point tensor parameter in the module tree.

fn freeze(self) -> Self

Freeze the whole module tree.

fn unfreeze(self) -> Self

Unfreeze the whole module tree.

fn freeze_group(self, group: ParamGroup) -> Self

Freeze every module-owned value in the given group, leaving the rest of the module untouched.

fn unfreeze_group(self, group: ParamGroup) -> Self

Unfreeze every module-owned value in the given group, leaving the rest of the module untouched.

fn num_params(&self) -> usize

Get the number of parameters the module has, including all of its sub-modules, without initializing the ones not initialized yet.

fn quantize_weights(self, quantizer: &mut Quantizer) -> Self

Quantize the weights of the module.

fn quantize_weights_group( self, quantizer: &mut Quantizer, group: ParamGroup, ) -> Self

Quantize the weights of the given parameter group.

fn apply_reparameterization<R>(self, reparameterizer: R) -> Self

where Self: Sized, R: Reparameterizer,

Attach reparameterizations using the given Reparameterizer.

fn apply_lora(self, lora: Lora) -> Self

where Self: Sized,

Attach LoRA adapters to the module’s 2-D weights, freezing all base tensor parameters.

fn apply_qlora(self, qlora: QLora) -> Self

where Self: Sized,

Apply QLoRA to the module: quantize the (frozen) base tensor parameters and attach trainable LoRA adapters to 2-D weights.

fn into_record(self) -> ModuleRecord

where Self: Sized,

Collect this module’s tensor parameters into a ModuleRecord.

fn into_record_group(self, group: ParamGroup) -> ModuleRecord

where Self: Sized,

Collect the tensor parameters group names into a ModuleRecord.

fn try_load_record(self, record: ModuleRecord) -> Result<Self, RecordError>

where Self: Sized,

Apply a ModuleRecord to this module, returning the loaded module.

fn load_record(self, record: ModuleRecord) -> Self

where Self: Sized,

Apply a ModuleRecord to this module, consuming and returning it.

fn save_file<P>(self, path: P) -> Result<(), RecordError>

where P: AsRef<Path>, Self: Sized,

Save this module’s parameters to a burnpack file on disk.

fn load_file<P>(self, path: P) -> Self

where P: AsRef<Path>, Self: Sized,

Load this module’s parameters from a burnpack file on disk, returning the loaded module.

fn try_load_file<P>(self, path: P) -> Result<Self, RecordError>

where P: AsRef<Path>, Self: Sized,

Fallible variant of load_file.

impl<const D: usize, K> ModuleDisplay for Tensor<D, K>

where K: Basic,

fn format(&self, passed_settings: DisplaySettings) -> String

Formats the module with provided display settings.

fn custom_settings(&self) -> Option<DisplaySettings>

Custom display settings for the module.

fn custom_content(&self, _content: Content) -> Option<Content>

Custom attributes for the module.

impl<const D: usize, K> ModuleDisplayDefault for Tensor<D, K>

where K: Basic,

fn content(&self, content: Content) -> Option<Content>

Attributes of the module used for display purposes.

fn num_params(&self) -> usize

Gets the number of the parameters of the module.

impl Morphology for Tensor<3>

fn erode(self, kernel: Tensor<2, Bool>, opts: MorphOptions) -> Tensor<3>

Erodes this tensor using the specified kernel. Assumes NHWC layout.

fn dilate(self, kernel: Tensor<2, Bool>, opts: MorphOptions) -> Tensor<3>

Dilates this tensor using the specified kernel. Assumes NHWC layout.

impl Morphology for Tensor<3, Int>

fn erode(self, kernel: Tensor<2, Bool>, opts: MorphOptions) -> Tensor<3, Int>

Erodes this tensor using the specified kernel. Assumes NHWC layout.

fn dilate(self, kernel: Tensor<2, Bool>, opts: MorphOptions) -> Tensor<3, Int>

Dilates this tensor using the specified kernel. Assumes NHWC layout.

impl Morphology for Tensor<3, Bool>

fn erode(self, kernel: Tensor<2, Bool>, opts: MorphOptions) -> Tensor<3, Bool>

Erodes this tensor using the specified kernel. Assumes NHWC layout.

fn dilate(self, kernel: Tensor<2, Bool>, opts: MorphOptions) -> Tensor<3, Bool>

Dilates this tensor using the specified kernel. Assumes NHWC layout.

impl<const D: usize, K> Mul for Tensor<D, K>

where K: Numeric,

type Output = Tensor<D, K>

The resulting type after applying the * operator.

fn mul(self, rhs: Tensor<D, K>) -> <Tensor<D, K> as Mul>::Output

Performs the * operation.

impl<E, const D: usize, K> Mul<E> for Tensor<D, K>

where E: ElementConversion, K: Numeric,

type Output = Tensor<D, K>

The resulting type after applying the * operator.

fn mul(self, other: E) -> <Tensor<D, K> as Mul<E>>::Output

Performs the * operation.

impl<const D: usize, K> Mul<Tensor<D, K>> for f64

where K: Numeric,

type Output = Tensor<D, K>

The resulting type after applying the * operator.

fn mul(self, other: Tensor<D, K>) -> <f64 as Mul<Tensor<D, K>>>::Output

Performs the * operation.

impl<const D: usize, K> Mul<Tensor<D, K>> for u32

where K: Numeric,

type Output = Tensor<D, K>

The resulting type after applying the * operator.

fn mul(self, other: Tensor<D, K>) -> <u32 as Mul<Tensor<D, K>>>::Output

Performs the * operation.

impl<const D: usize, K> Neg for Tensor<D, K>

where K: Numeric,

type Output = Tensor<D, K>

The resulting type after applying the - operator.

fn neg(self) -> <Tensor<D, K> as Neg>::Output

Performs the unary - operation.

impl Nms for Tensor<2>

fn nms(self, scores: Tensor<1>, options: NmsOptions) -> Tensor<1, Int>

Perform Non-Maximum Suppression on this tensor of bounding boxes.

impl<const D: usize> Not for Tensor<D, Bool>

type Output = Tensor<D, Bool>

The resulting type after applying the ! operator.

fn not(self) -> <Tensor<D, Bool> as Not>::Output

Performs the unary ! operation.

impl<const D: usize> Parameter for Tensor<D>

fn is_require_grad(&self) -> bool

Fetch the gradient requirement.

fn set_require_grad(self, require_grad: bool) -> Tensor<D>

Set the gradient requirement.

fn device(&self) -> Device

Fetch the device.

fn shape(&self) -> Shape

Fetch the shape of the parameter.

fn load_to_device(self, device: &Device) -> Tensor<D>

Moves the parameter to the target device if it is not already on it, applying any kind-specific preparation required for the loading lifecycle (e.g. detach).

impl<const D: usize> Parameter for Tensor<D, Int>

fn is_require_grad(&self) -> bool

Fetch the gradient requirement.

fn set_require_grad(self, _require_grad: bool) -> Tensor<D, Int>

Set the gradient requirement.

fn device(&self) -> Device

Fetch the device.

fn shape(&self) -> Shape

Fetch the shape of the parameter.

fn load_to_device(self, device: &Device) -> Tensor<D, Int>

Moves the parameter to the target device if it is not already on it, applying any kind-specific preparation required for the loading lifecycle (e.g. detach).

impl<const D: usize> Parameter for Tensor<D, Bool>

fn is_require_grad(&self) -> bool

Fetch the gradient requirement.

fn set_require_grad(self, _require_grad: bool) -> Tensor<D, Bool>

Set the gradient requirement.

fn device(&self) -> Device

Fetch the device.

fn shape(&self) -> Shape

Fetch the shape of the parameter.

fn load_to_device(self, device: &Device) -> Tensor<D, Bool>

Moves the parameter to the target device if it is not already on it, applying any kind-specific preparation required for the loading lifecycle (e.g. detach).

impl<const D: usize> ParameterValue for Tensor<D>

impl<const D: usize> ParameterValue for Tensor<D, Int>

impl<const D: usize> ParameterValue for Tensor<D, Bool>

impl<const D: usize, K> Rem for Tensor<D, K>

where K: Numeric,

type Output = Tensor<D, K>

The resulting type after applying the % operator.

fn rem(self, rhs: Tensor<D, K>) -> <Tensor<D, K> as Rem>::Output

Performs the % operation.

impl<E, const D: usize, K> Rem<E> for Tensor<D, K>

where E: ElementConversion, K: Numeric,

type Output = Tensor<D, K>

The resulting type after applying the % operator.

fn rem(self, other: E) -> <Tensor<D, K> as Rem<E>>::Output

Performs the % operation.

impl<const D: usize, K> Serialize for Tensor<D, K>

where K: Basic,

fn serialize<S>( &self, serializer: S, ) -> Result<<S as Serializer>::Ok, <S as Serializer>::Error>

where S: Serializer,

Serialize this value into the given Serde serializer.

impl SliceAccess for Tensor<2>

fn zeros_like(sample: &Tensor<2>, capacity: usize, device: &Device) -> Tensor<2>

Create zeroed storage matching the shape of sample but with capacity rows along the first dimension.

fn select(self, dim: usize, indices: Tensor<1, Int>) -> Tensor<2>

Select rows at the given indices along the specified dimension.

fn slice_assign_inplace(&mut self, index: usize, value: Tensor<2>)

Assign value at row index along the first dimension, in place.

impl<const D: usize, K> Sub for Tensor<D, K>

where K: Numeric,

type Output = Tensor<D, K>

The resulting type after applying the - operator.

fn sub(self, rhs: Tensor<D, K>) -> <Tensor<D, K> as Sub>::Output

Performs the - operation.

impl<E, const D: usize, K> Sub<E> for Tensor<D, K>

where E: ElementConversion, K: Numeric,

type Output = Tensor<D, K>

The resulting type after applying the - operator.

fn sub(self, other: E) -> <Tensor<D, K> as Sub<E>>::Output

Performs the - operation.

impl<const D: usize, K> Sub<Tensor<D, K>> for f64

where K: Numeric,

type Output = Tensor<D, K>

The resulting type after applying the - operator.

fn sub(self, tensor: Tensor<D, K>) -> <f64 as Sub<Tensor<D, K>>>::Output

Performs the - operation.

impl<const D: usize, K> Sub<Tensor<D, K>> for u32

where K: Numeric,

type Output = Tensor<D, K>

The resulting type after applying the - operator.

fn sub(self, tensor: Tensor<D, K>) -> <u32 as Sub<Tensor<D, K>>>::Output

Performs the - operation.