Struct burn_tensor::Tensor

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

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

Implementations

impl<const D: usize, B: AutodiffBackend> Tensor<B, D>

pub fn backward(&self) -> B::Gradients

Backward pass of the tensor.

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

Get the gradients of a tensor if it exist.

Returns a new reference to the same tensor. Therefore the same grad tensor can be accessed multiple times. If you only need to get the gradients one time, consider using grad_remove for better performance.

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

Remove the grad tensor from the grads struct returning the result.

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

Replace the grad tensor from the grads struct with the provided gradient.

impl<const D: usize, B: AutodiffBackend, K: BasicAutodiffOps<B>> Tensor<B, D, K>

pub fn inner(self) -> Tensor<B::InnerBackend, D, K::InnerKind>

Returns the inner tensor without the autodiff information.

pub fn from_inner(inner: Tensor<B::InnerBackend, D, K::InnerKind>) -> Self

Convert a tensor to the autodiff backend.

Arguments

• inner - The tensor to convert.

Returns

The tensor converted to the autodiff backend.

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

where B: Backend, K: TensorKind<B>,

pub fn new(primitive: K::Primitive<D>) -> Self

Constructs a new Tensor.

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

where B: Backend, K: BasicOps<B>,

pub fn into_primitive(self) -> K::Primitive<D>

Converts the tensor into a primitive tensor.

pub fn from_primitive(tensor: K::Primitive<D>) -> Self

Converts from a primitive tensor into a tensor.

pub fn empty<S: Into<Shape<D>>>(shape: S, device: &B::Device) -> Self

Create an empty tensor of the given shape.

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

Returns the dimensions of the current tensor.

Equivalent to tensor.shape().dims.

pub fn shape(&self) -> Shape<D>

Returns the shape of the current tensor.

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

Reshape the tensor to have the given shape.

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::backend::Backend;
use burn_tensor::Tensor;

fn example<B: Backend>() {
   let device = Default::default();
   let tensor = Tensor::<B, 3>::ones([2, 3, 4], &device);
   // Given a 3D tensor with dimensions (2, 3, 4), reshape it to (2, 12)
   let reshaped_tensor: Tensor::<B, 2> = tensor.reshape([2, -1]);
   // The resulting tensor will have dimensions (2, 12).
   println!("{:?}", reshaped_tensor.shape());
}

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

Transpose the tensor.

Arguments

• tensor - The tensor to transpose.

Returns

The transposed tensor.

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

Swaps two dimensions of a tensor.

Arguments

• tensor - The tensor to swap the dimensions of.
• dim1 - The first dimension to swap.
• dim2 - The second dimension to swap.

Returns

The tensor with the dimensions swapped.

pub fn flatten<const D2: usize>( self, start_dim: usize, end_dim: usize ) -> Tensor<B, 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.
• end_dim: The ending dimension of the range to be flattened (inclusive).

Type Parameters

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

Returns

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

Example

use burn_tensor::backend::Backend;
use burn_tensor::{Tensor, Shape};

fn example<B: Backend>() {
    let device = Default::default();
    let tensor = Tensor::<B, 3>::ones(Shape::new([2, 3, 4]), &device);

    // Given a 3D tensor with dimensions (2, 3, 4), flatten the dimensions between indices 1 and 2:
    let flattened_tensor: Tensor::<B, 2> = tensor.flatten(1, 2);

    // The resulting tensor will have dimensions (2, 12).
   println!("{:?}", flattened_tensor.shape());
}

pub fn squeeze<const D2: usize>(self, dim: usize) -> Tensor<B, 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.

Type Parameters

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

Returns

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

Example

use burn_tensor::backend::Backend;
use burn_tensor::{Tensor, Shape};

fn example<B: Backend>() {
    let device = Default::default();
    let tensor = Tensor::<B, 3>::ones(Shape::new([2, 1, 4]), &device);

    // Given a 3D tensor with dimensions (2, 1, 4), squeeze the dimension 1
    let squeezed_tensor: Tensor::<B, 2> = tensor.squeeze(1);

    // Resulting tensor will have dimensions (2, 4)
    println!("{:?}", squeezed_tensor.shape());
}

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

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

If the output size is higher than the current tensor.

Example

use burn_tensor::backend::Backend;
use burn_tensor::{Tensor, Shape};

fn example<B: Backend>() {
    let device = Default::default();
    let tensor = Tensor::<B, 2>::ones(Shape::new([3, 3]), &device);
    let tensor = tensor.unsqueeze::<4>();
    println!("{:?}", tensor.shape());
    // Shape { dims: [1, 1, 3, 3] }
}

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

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

Example

use burn_tensor::backend::Backend;
use burn_tensor::{Tensor, Shape};

fn example<B: Backend>() {
    let device = Default::default();
    let tensor = Tensor::<B, 2>::ones(Shape::new([3, 3]), &device);
    let tensor: Tensor<B, 3> = tensor.unsqueeze_dim(1);
    println!("{:?}", tensor.shape());
    // Shape { dims: [3, 1, 3] }
}

pub fn slice<const D2: usize>(self, ranges: [Range<usize>; D2]) -> Self

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

Panics

If a range exceeds the number of elements on a dimension.

Example

use burn_tensor::backend::Backend;
use burn_tensor::{Tensor, Shape};

fn example<B: Backend>() {
    let device = B::Device::default();
    // Create a tensor with a single dimension of ints between 0 and 11
    let tensor = Tensor::<B, 1, burn_tensor::Int>::arange(0..12, &device);
    // Select elements 0, 1, 2, 3 from the first dimension
    let tensor_slices = tensor.clone().slice([0..4]);
    println!("\nexpecting [0,1,2,3] : {:?}", tensor);
    println!("expecting [4] : {:?}", tensor.dims());

    // Create a Tensor with 3 dimensions
    let tensor = Tensor::<B, 3>::ones(Shape::new([2, 3, 3]), &device);
    // This slice will select the element 0 on the first dimension,
    // elements 0,1,2 of the second dimension and element 1 of third dimension
    let tensor_slices = tensor.slice([0..1, 0..3, 1..2]);
    println!("expecting [1, 3, 1] : {:?}", tensor_slices.dims());

    // Create a tensor of ints from 0 to 11 and reshape it into three dimensions
    let tensor = Tensor::<B, 1, burn_tensor::Int>::arange(0..12, &device);
    let tensor = tensor.reshape([1, 3, 4]);
    println!("\nexpecting [[[0,1,2,3],[4,5,6,7],[8,9,10,11]]] : {:?}", tensor);
    println!("expecting [1, 3, 4] : {:?}", tensor.dims());
    // Select element 0 of first dimension, elements 1,2 of second dimension
    // and element 1 of third dimension
    //
    // This is the equivalent of this pseudo code
    // let mut v = vec![[[]]];
    // v[0][0][0] = tensor[0][1][1];
    // v[0][1][0] = tensor[0][2][1];
    let tensor_slices = tensor.slice([0..1, 1..3, 1..2]);
    println!("\nexpecting [1, 2, 1] : {:?}", tensor_slices.dims());
    println!("expecting [[[5],[9]]] : {:?}", tensor_slices);
}

pub fn slice_assign<const D2: usize>( self, ranges: [Range<usize>; D2], values: Self ) -> Self

Returns a copy of the current tensor with the selected elements changed to the new ones at the selected indices.

Panics

• If a range exceeds the number of elements on a dimension.
• If the given values don’t match the given ranges.

Example

use burn_tensor::backend::Backend;
use burn_tensor::Tensor;

fn example<B: Backend>() {
    let device = B::Device::default();
    let tensor = Tensor::<B, 3>::ones([2, 3, 3], &device);
    let values = Tensor::<B, 3>::zeros([1, 1, 1], &device);
    let tensor_sliced = tensor.slice_assign([0..1, 0..1, 0..1], values);
    println!("{:?}", tensor_sliced.dims()); // [2, 3, 3]
}

pub fn device(&self) -> B::Device

Returns the device of the current tensor.

pub fn to_device(self, device: &B::Device) -> Self

Returns a new tensor on the given device.

pub fn into_data(self) -> Data<K::Elem, D>

Returns the data of the current tensor.

pub fn to_data(&self) -> Data<K::Elem, D>

Returns the data of the current tensor without taking ownership.

pub fn from_data<T>(data: T, device: &B::Device) -> Self

where T: Into<Data<K::Elem, D>>,

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

pub fn repeat(self, dim: usize, times: usize) -> Self

Repeat the tensor along the given dimension.

Panics

If the selected dimension more than one item.

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

Applies element wise equal comparison and returns a boolean tensor.

Panics

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

pub fn cat(tensors: Vec<Self>, dim: usize) -> Self

Concatenates all tensors into a new one along the given dimension.

Panics

If all tensors don’t have the same shape.

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

Concatenates all tensors into a new one along a new dimension.

Panics

If all tensors don’t have the same shape. Given dimension is not with range of 0..D2

pub fn iter_dim(self, dim: usize) -> DimIter<B, D, K>

Iterate over slices of tensors alongside a given dimension.

Panics

Given dimension is less than tensor rank.

Returns

A tensor iterator.

pub fn narrow(self, dim: usize, start: usize, length: usize) -> Self

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

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.

pub fn chunk(self, chunks: usize, dim: usize) -> Vec<Self>

Attempts to split the tensor along the given dimension into chunks. 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.

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

where B: Backend,

pub fn from_bool(data: Data<bool, D>, device: &B::Device) -> Self

Create a boolean tensor from data on the given device.

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

Convert the bool tensor into an int tensor.

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

Convert the bool tensor into an float tensor.

pub fn bool_not(self) -> Self

Inverses boolean values.

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

where B: Backend,

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

Executes an operation on the tensor and modifies its value.

Notes

This won’t necessary 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 exp(self) -> Self

Applies element wise exponential operation.

y = e^x

pub fn log(self) -> Self

Applies element wise natural log operation ln.

y = log(x)

pub fn log1p(self) -> Self

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

y = log(x+1)

pub fn erf(self) -> Self

Applies the error function element wise.

y = erf(x)

pub fn recip(self) -> Self

Applies element wise reciprocal operation.

pub fn sqrt(self) -> Self

Applies element wise root square operation.

pub fn cos(self) -> Self

Applies element wise cosine operation.

pub fn sin(self) -> Self

Applies element wise sine operation.

pub fn tanh(self) -> Self

Applies element wise hyperbolic tangent operation.

pub fn from_floats<A: Into<Data<f32, D>>>(floats: A, device: &B::Device) -> Self

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

Example

use burn_tensor::backend::Backend;
use burn_tensor::Tensor;

fn example<B: Backend>() {
    let device = B::Device::default();
    let _ = Tensor::<B, 1>::from_floats([1.0, 2.0], &device);
    let _ = Tensor::<B, 2>::from_floats([[1.0, 2.0], [3.0, 4.0]], &device);
}

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

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

Example

use burn_tensor::backend::Backend;
use burn_tensor::Tensor;

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

pub fn zeros_like(&self) -> Self

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

pub fn ones_like(&self) -> Self

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

pub fn random_like(&self, distribution: Distribution) -> Self

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

pub fn one_hot(index: usize, num_classes: usize, device: &B::Device) -> Self

Create a one hot tensor.

Example

use burn_tensor::backend::Backend;
use burn_tensor::Tensor;

fn example<B: Backend>() {
    let device = Default::default();
    let one_hot = Tensor::<B, 1>::one_hot(2, 10, &device);
    println!("{}", one_hot.to_data());
    // [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
}

pub fn matmul(self, other: Self) -> Self

Applies the matrix multiplication operation.

C = AB

Panics

If the two tensors dont’ have a compatible shape.

pub fn var(self, dim: usize) -> Self

Calculate the variance along the given dimension.

pub fn var_bias(self, dim: usize) -> Self

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

pub fn var_mean(self, dim: usize) -> (Self, Self)

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

pub fn var_mean_bias(self, dim: usize) -> (Self, Self)

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

pub fn random<S: Into<Shape<D>>>( shape: S, distribution: Distribution, device: &B::Device ) -> Self

Create a random tensor of the given shape on the given device where each element is sampled from the given distribution.

pub fn to_full_precision(&self) -> Tensor<B::FullPrecisionBackend, D>

Returns a tensor with full precision based on the selected backend.

pub fn from_full_precision(tensor: Tensor<B::FullPrecisionBackend, D>) -> Self

Returns a tensor on the selected backend from a full precision tensor.

pub fn detach(self) -> Self

Detach the current tensor from the autodiff graph.

This function does nothing when autodiff is not enabled. This can be used in batchers or elsewhere to ensure that previous operations are not considered in the autodiff graph.

pub fn require_grad(self) -> Self

Mark the tensor to keep gradients during the backward pass.

This function does nothing when autodiff is not enabled.

pub fn is_require_grad(&self) -> bool

Returns true if the tensor requires gradients during the backward pass.

pub fn set_require_grad(self, require_grad: bool) -> Self

Mark the tensor as tracked or untracked depending on the require grad argument. When tracked, the gradients will be available after the backward pass.

This function does nothing when autodiff is not enabled.

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

Calculate covaraince matrix between different entries alongside a given dimension.

Arguments

• size - The size of the square matrix.
• correction_factor - Is usually 1 for samples and 0 for population.

impl<B> Tensor<B, 1, Int>

where B: Backend,

pub fn arange(range: Range<usize>, device: &B::Device) -> Self

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<usize>, step: usize, device: &B::Device) -> Self

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, B> Tensor<B, D, Int>

where B: Backend,

pub fn from_ints<A: Into<Data<i32, D>>>(ints: A, device: &B::Device) -> Self

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

Example

use burn_tensor::backend::Backend;
use burn_tensor::{Tensor, Int};

fn example<B: Backend>() {
    let device = B::Device::default();
    let _x: Tensor<B, 1, Int> = Tensor::from_ints([1, 2], &device);
    let _y: Tensor<B, 2, Int> = Tensor::from_ints([[1, 2], [3, 4]], &device);
}

pub fn float(self) -> Tensor<B, D, Float>

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

Example

use burn_tensor::backend::Backend;
use burn_tensor::{Int, Tensor};

fn example<B: Backend>() {
    let device = Default::default();
    let int_tensor = Tensor::<B, 1, Int>::arange(0..5, &device);
    let float_tensor = int_tensor.float();
}

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

where B: Backend, K: Numeric<B>, K::Elem: Element,

pub fn into_scalar(self) -> K::Elem

Convert the tensor into a scalar.

Panics

If the tensor doesn’t have one element.

pub fn add(self, other: Self) -> Self

Applies element wise addition operation.

y = x2 + x1

pub fn add_scalar<E: ElementConversion>(self, other: E) -> Self

Applies element wise addition operation with a scalar.

y = x + s

pub fn sub(self, other: Self) -> Self

Applies element wise subtraction operation.

y = x2 - x1

pub fn sub_scalar<E: ElementConversion>(self, other: E) -> Self

Applies element wise subtraction operation with a scalar.

y = x - s

pub fn div(self, other: Self) -> Self

Applies element wise division operation.

y = x2 / x1

pub fn div_scalar<E: ElementConversion>(self, other: E) -> Self

Applies element wise division operation with a scalar.

y = x / s

pub fn mul(self, other: Self) -> Self

Applies element wise multiplication operation.

y = x2 * x1

pub fn mul_scalar<E: ElementConversion>(self, other: E) -> Self

Applies element wise multiplication operation with a scalar.

y = x * s

pub fn neg(self) -> Self

Switch sign of each element in the tensor.

y = -x

pub fn zeros<S: Into<Shape<D>>>(shape: S, device: &B::Device) -> Self

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

pub fn ones<S: Into<Shape<D>>>(shape: S, device: &B::Device) -> Self

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

pub fn full<S: Into<Shape<D>>, E: ElementConversion>( shape: S, fill_value: E, device: &B::Device ) -> Self

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

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

Aggregate all elements in the tensor with the mean operation.

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

Aggregate all elements in the tensor with the sum operation.

pub fn mean_dim(self, dim: usize) -> Self

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

pub fn sum_dim(self, dim: usize) -> Self

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

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

Applies element wise equal comparison and returns a boolean tensor.

pub fn greater(self, other: Self) -> Tensor<B, D, Bool>

Applies element wise greater comparison and returns a boolean tensor.

Panics

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

pub fn greater_equal(self, other: Self) -> Tensor<B, D, Bool>

Applies element wise greater-equal comparison and returns a boolean tensor.

Panics

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

pub fn lower(self, other: Self) -> Tensor<B, D, Bool>

Applies element wise lower comparison and returns a boolean tensor.

Panics

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

pub fn lower_equal(self, other: Self) -> Tensor<B, D, Bool>

Applies element wise lower-equal comparison and returns a boolean tensor.

Panics

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

pub fn greater_elem<E: ElementConversion>(self, other: E) -> Tensor<B, D, Bool>

Applies element wise greater comparison and returns a boolean tensor.

pub fn greater_equal_elem<E: ElementConversion>( self, other: E ) -> Tensor<B, D, Bool>

Applies element wise greater-equal comparison and returns a boolean tensor.

pub fn lower_elem<E: ElementConversion>(self, other: E) -> Tensor<B, D, Bool>

Applies element wise lower comparison and returns a boolean tensor.

pub fn lower_equal_elem<E: ElementConversion>( self, other: E ) -> Tensor<B, D, Bool>

Applies element wise lower-equal comparison and returns a boolean tensor.

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

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.

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

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.

pub fn gather(self, dim: usize, indices: Tensor<B, D, Int>) -> Self

Gather tensor elements corresponding to the given indices from the specified dim.

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.

pub fn scatter( self, dim: usize, indices: Tensor<B, D, Int>, values: Self ) -> Self

Assign the gathered elements corresponding to the given indices along the specified dimension from the value tensor to the original tensor using sum reduction.

Example using a 3D tensor:

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

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.

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

pub fn select(self, dim: usize, indices: Tensor<B, 1, Int>) -> Self

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

Example using a 3D tensor:

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

pub fn select_assign( self, dim: usize, indices: Tensor<B, 1, Int>, values: Tensor<B, D, K> ) -> Self

Assign the selected elements along the given dimension corresponding to the given indices from the value tensor to the original tensor using sum reduction.

Example using a 3D tensor:

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

pub fn argmax(self, dim: usize) -> Tensor<B, D, Int>

Applies the argmax function along the given dimension and returns an integer tensor.

Example

use burn_tensor::backend::Backend;
use burn_tensor::{Tensor, Shape};

fn example<B: Backend>() {
    let device = B::Device::default();
    let tensor = Tensor::<B, 3>::ones(Shape::new([2, 3, 3]), &device);
    let tensor = tensor.argmax(1);
    println!("{:?}", tensor.shape());
    // Shape { dims: [2, 1, 3] }
}

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

Find the maximum value.

pub fn max_dim(self, dim: usize) -> Tensor<B, D, K>

Find the maximum value along the given dimension.

pub fn max_dim_with_indices( self, dim: usize ) -> (Tensor<B, D, K>, Tensor<B, D, Int>)

Find the maximum value along the given dimension.

Also returns the indices.

pub fn argmin(self, dim: usize) -> Tensor<B, D, Int>

Applies the argmin function along the given dimension and returns an integer tensor.

Example

use burn_tensor::backend::Backend;
use burn_tensor::{Tensor, Shape};

fn example<B: Backend>() {
    let device = Default::default();
    let tensor = Tensor::<B, 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<B, 1, K>

Find the minimum value.

pub fn min_dim(self, dim: usize) -> Tensor<B, D, K>

Find the minimum value along the given dimension.

pub fn min_dim_with_indices( self, dim: usize ) -> (Tensor<B, D, K>, Tensor<B, D, Int>)

Find the minimum value along the given dimension.

Also returns the indices.

pub fn clamp<E: ElementConversion>(self, min: E, max: E) -> Self

Clamp the tensor between the given min and max values.

Arguments

• min - The minimum value.
• max - The maximum value.

Returns

A new tensor with the values clamped between the given min and max values.

pub fn clamp_min<E: ElementConversion>(self, min: E) -> Self

Clamps a tensor 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.

pub fn clamp_max<E: ElementConversion>(self, max: E) -> Self

Clamps a tensor 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.

pub fn abs(self) -> Self

Apply element wise absolute value operation

pub fn triu(self, diagonal: i64) -> Self

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.

Example

use burn_tensor::backend::Backend;
use burn_tensor::{Int, Tensor};

fn example<B: Backend>() {
   let device = Default::default();
   let tensor = Tensor::<B, 2, Int>::from_ints(
       [
         [1, 2, 3],
         [4, 5, 6],
         [7, 8, 9]
       ],
       &device
   );
   let tensor = tensor.triu(1);
   println!("{}", tensor);
   // Tensor { data: [
   //   [0, 2, 3],
   //   [0, 0, 6],
   //   [0, 0, 0]
   // ], ... }
}

pub fn tril(self, diagonal: i64) -> Self

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.

Example

use burn_tensor::backend::Backend;
use burn_tensor::{Int, Tensor};

fn example<B: Backend>() {
   let device = Default::default();
   let tensor = Tensor::<B, 2, Int>::from_ints(
       [
         [1, 2, 3],
         [4, 5, 6],
         [7, 8, 9]
       ],
       &device
   );

   let tensor = tensor.tril(-1);
   println!("{}", tensor);
   // Tensor { data: [
   //   [0, 0, 0],
   //   [4, 0, 0],
   //   [7, 8, 0]
   // ], ... }
}

pub fn powf(self, other: Self) -> Self

Applies element wise power operation with a float Tensor

pub fn powf_scalar<E: ElementConversion>(self, other: E) -> Self

Applies element wise power operation with a float scalar

pub fn powi(self, other: Self) -> Self

Applies element wise power operation with a integer Tensor

pub fn powi_scalar<E: ElementConversion>(self, other: E) -> Self

Applies element wise power operation with a integer scalar

impl<B, K> Tensor<B, 2, K>

where B: Backend, K: Numeric<B>, K::Elem: Element,

pub fn diagonal(size: usize, device: &B::Device) -> Self

Create diagonal matrix.

Arguments

• size - The size of the square matrix.

Trait Implementations

impl<E, const D: usize, B, K> Add<E> for Tensor<B, D, K>

where E: ElementConversion, B: Backend, K: Numeric<B>, K::Elem: Element,

type Output = Tensor<B, D, K>

The resulting type after applying the + operator.

fn add(self, other: E) -> Self

Performs the + operation.

impl<B, const D: usize, K> Add for Tensor<B, D, K>

where B: Backend, K: Numeric<B>, K::Elem: Element,

type Output = Tensor<B, D, K>

The resulting type after applying the + operator.

fn add(self, rhs: Tensor<B, D, K>) -> Self

Performs the + operation.

impl<B: Backend, const D: usize> BitXor<T> for Tensor<B, D>

type Output = Tensor<B, D>

The resulting type after applying the ^ operator.

fn bitxor(self, _: T) -> Self::Output

Performs the ^ operation.

impl<B, const D: usize, K> Clone for Tensor<B, D, K>

where B: Backend + Clone, K: TensorKind<B> + Clone, K::Primitive<D>: Clone,

fn clone(&self) -> Tensor<B, D, K>

Returns a copy of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl<B, const D: usize, K> Debug for Tensor<B, D, K>

where B: Backend + Debug, K: TensorKind<B> + Debug, K::Primitive<D>: Debug,

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl<'de, B, const D: usize, K> Deserialize<'de> for Tensor<B, D, K>

where B: Backend, K: BasicOps<B>, K::Elem: Debug + Copy + Deserialize<'de>,

fn deserialize<De: Deserializer<'de>>( deserializer: De ) -> Result<Self, De::Error>

Deserialize this value from the given Serde deserializer.

impl<B, const D: usize, K> Display for Tensor<B, D, K>

where B: Backend, B::IntElem: Display, K: BasicOps<B>, <K as BasicOps<B>>::Elem: Debug,

Pretty print tensors

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl<E, const D: usize, B, K> Div<E> for Tensor<B, D, K>

where E: ElementConversion, B: Backend, K: Numeric<B>, K::Elem: Element,

type Output = Tensor<B, D, K>

The resulting type after applying the / operator.

fn div(self, other: E) -> Self

Performs the / operation.

impl<B, const D: usize, K> Div for Tensor<B, D, K>

where B: Backend, K: Numeric<B>, K::Elem: Element,

type Output = Tensor<B, D, K>

The resulting type after applying the / operator.

fn div(self, rhs: Tensor<B, D, K>) -> Self

Performs the / operation.

impl<B, const D: usize, K, T> From<T> for Tensor<B, D, K>

where B: Backend, K: BasicOps<B>, T: Into<Data<K::Elem, D>>,

fn from(value: T) -> Self

Converts to this type from the input type.

impl<E, const D: usize, B, K> Mul<E> for Tensor<B, D, K>

where E: ElementConversion, B: Backend, K: Numeric<B>, K::Elem: Element,

type Output = Tensor<B, D, K>

The resulting type after applying the * operator.

fn mul(self, other: E) -> Self

Performs the * operation.

impl<B, const D: usize, K> Mul for Tensor<B, D, K>

where B: Backend, K: Numeric<B>, K::Elem: Element,

type Output = Tensor<B, D, K>

The resulting type after applying the * operator.

fn mul(self, rhs: Tensor<B, D, K>) -> Self

Performs the * operation.

impl<B, const D: usize, K> Neg for Tensor<B, D, K>

where B: Backend, K: Numeric<B>, K::Elem: Element,

type Output = Tensor<B, D, K>

The resulting type after applying the - operator.

fn neg(self) -> Self

Performs the unary - operation.

impl<B, const D: usize, K> Serialize for Tensor<B, D, K>

where B: Backend, K: BasicOps<B>, K::Elem: Debug + Copy + Serialize,

fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error>

Serialize this value into the given Serde serializer.

impl<E, const D: usize, B, K> Sub<E> for Tensor<B, D, K>

where E: ElementConversion, B: Backend, K: Numeric<B>, K::Elem: Element,

type Output = Tensor<B, D, K>

The resulting type after applying the - operator.

fn sub(self, other: E) -> Self

Performs the - operation.

impl<B, const D: usize, K> Sub for Tensor<B, D, K>

where B: Backend, K: Numeric<B>, K::Elem: Element,

type Output = Tensor<B, D, K>

The resulting type after applying the - operator.

fn sub(self, rhs: Tensor<B, D, K>) -> Self

Performs the - operation.