Struct burn_tensor::Tensor
source · pub struct Tensor<B, const D: usize, K = Float>where
B: Backend,
K: TensorKind<B>,{ /* private fields */ }
Expand description
A tensor with a given backend, shape and data type.
Implementations§
source§impl<B, const D: usize, K> Tensor<B, D, K>where
B: Backend,
K: BasicOps<B>,
impl<B, const D: usize, K> Tensor<B, D, K>where B: Backend, K: BasicOps<B>,
sourcepub fn into_primitive(self) -> K::Primitive<D>
pub fn into_primitive(self) -> K::Primitive<D>
Converts the tensor into a primitive tensor.
sourcepub fn from_primitive(tensor: K::Primitive<D>) -> Self
pub fn from_primitive(tensor: K::Primitive<D>) -> Self
Converts from a primitive tensor into a tensor.
sourcepub fn empty_device<S: Into<Shape<D>>>(shape: S, device: &B::Device) -> Self
pub fn empty_device<S: Into<Shape<D>>>(shape: S, device: &B::Device) -> Self
Create an empty tensor of the given shape.
sourcepub fn dims(&self) -> [usize; D]
pub fn dims(&self) -> [usize; D]
Returns the dimensions of the current tensor.
Equivalent to tensor.shape().dims
.
sourcepub fn reshape<const D2: usize, S: ReshapeArgs<D2>>(
self,
shape: S
) -> Tensor<B, D2, K>
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 tensor = Tensor::<B, 3>::ones([2, 3, 4]);
// 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());
}
sourcepub fn flatten<const D2: usize>(
self,
start_dim: usize,
end_dim: usize
) -> Tensor<B, D2, K>
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 tensor = Tensor::<B, 3>::ones(Shape::new([2, 3, 4]));
// 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());
}
sourcepub fn squeeze<const D2: usize>(self, dim: usize) -> Tensor<B, D2, K>
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 tensor = Tensor::<B, 3>::ones(Shape::new([2, 1, 4]));
// 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());
}
sourcepub fn unsqueeze<const D2: usize>(self) -> Tensor<B, D2, K>
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 tensor = Tensor::<B, 2>::ones(Shape::new([3, 3]));
let tensor = tensor.unsqueeze::<4>();
println!("{:?}", tensor.shape());
// Shape { dims: [1, 1, 3, 3] }
}
sourcepub fn slice<const D2: usize>(self, ranges: [Range<usize>; D2]) -> Self
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 tensor = Tensor::<B, 3>::ones(Shape::new([2, 3, 3]));
let tensor_slices = tensor.slice([0..1, 0..3, 1..2]);
println!("{:?}", tensor_slices.dims()); // [1, 3, 2]
}
sourcepub fn slice_assign<const D2: usize>(
self,
ranges: [Range<usize>; D2],
values: Self
) -> Self
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 tensor = Tensor::<B, 3>::ones([2, 3, 3]);
let values = Tensor::<B, 3>::zeros([1, 1, 1]);
let tensor_sliced = tensor.slice_assign([0..1, 0..1, 0..1], values);
println!("{:?}", tensor_sliced.dims()); // [2, 3, 3]
}
sourcepub fn to_data(&self) -> Data<K::Elem, D>
pub fn to_data(&self) -> Data<K::Elem, D>
Returns the data of the current tensor without taking ownership.
sourcepub fn from_data<T>(data: T) -> Selfwhere
T: Into<Data<K::Elem, D>>,
pub fn from_data<T>(data: T) -> Selfwhere T: Into<Data<K::Elem, D>>,
Create a tensor from the given data.
sourcepub fn from_data_device<T>(data: T, device: &B::Device) -> Selfwhere
T: Into<Data<K::Elem, D>>,
pub fn from_data_device<T>(data: T, device: &B::Device) -> Selfwhere T: Into<Data<K::Elem, D>>,
Create a tensor from the given data on the given device.
sourcepub fn equal(self, other: Self) -> Tensor<B, D, Bool>
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.
source§impl<B, const D: usize> Tensor<B, D, Bool>where
B: Backend,
impl<B, const D: usize> Tensor<B, D, Bool>where B: Backend,
sourcepub fn from_bool_device(data: Data<bool, D>, device: &B::Device) -> Self
pub fn from_bool_device(data: Data<bool, D>, device: &B::Device) -> Self
Create a boolean tensor from data on the given device.
source§impl<const D: usize, B> Tensor<B, D>where
B: Backend,
impl<const D: usize, B> Tensor<B, D>where B: Backend,
sourcepub fn inplace<F: FnOnce(Self) -> Self>(&mut self, func: F)
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.
sourcepub fn log1p(self) -> Self
pub fn log1p(self) -> Self
Applies the natural logarithm of one plus the input tensor, element-wise.
y = log(x+1)
sourcepub fn erf(self) -> Self
pub fn erf(self) -> Self
Applies the error function element wise.
y = erf(x)
sourcepub fn from_floats<A: Into<Data<f32, D>>>(floats: A) -> Self
pub fn from_floats<A: Into<Data<f32, D>>>(floats: A) -> Self
Create a tensor from floats (f32).
Example
use burn_tensor::backend::Backend;
use burn_tensor::Tensor;
fn example<B: Backend>() {
let _ = Tensor::<B, 1>::from_floats([1.0, 2.0]);
let _ = Tensor::<B, 2>::from_floats([[1.0, 2.0], [3.0, 4.0]]);
}
sourcepub fn int(self) -> Tensor<B, D, Int>
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 float_tensor = Tensor::<B, 1>::from_floats([1.0, 2.0]);
let int_tensor = float_tensor.int();
}
sourcepub fn zeros_like(&self) -> Self
pub fn zeros_like(&self) -> Self
Returns a new tensor with the same shape and device as the current tensor filled with zeros.
sourcepub fn ones_like(&self) -> Self
pub fn ones_like(&self) -> Self
Returns a new tensor with the same shape and device as the current tensor filled with ones.
sourcepub fn random_like(&self, distribution: Distribution<B::FloatElem>) -> Self
pub fn random_like(&self, distribution: Distribution<B::FloatElem>) -> Self
Returns a new tensor with the same shape and device as the current tensor filled random values sampled from the given distribution.
sourcepub fn one_hot(index: usize, num_classes: usize) -> Self
pub fn one_hot(index: usize, num_classes: usize) -> Self
Create a one hot tensor.
Example
use burn_tensor::backend::Backend;
use burn_tensor::Tensor;
fn example<B: Backend>() {
let one_hot = Tensor::<B, 1>::one_hot(2, 10);
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]
}
sourcepub fn matmul(self, other: Self) -> Self
pub fn matmul(self, other: Self) -> Self
Applies the matrix multiplication operation.
C = AB
Panics
If the two tensors dont’ have a compatible shape.
sourcepub fn var_bias(self, dim: usize) -> Self
pub fn var_bias(self, dim: usize) -> Self
Calculate the variance along the given dimension without applying the Bessel’s correction.
sourcepub fn var_mean(self, dim: usize) -> (Self, Self)
pub fn var_mean(self, dim: usize) -> (Self, Self)
Calculate the variance along the given dimension and also returns the mean.
sourcepub fn var_mean_bias(self, dim: usize) -> (Self, Self)
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.
sourcepub fn random<S: Into<Shape<D>>>(
shape: S,
distribution: Distribution<B::FloatElem>
) -> Self
pub fn random<S: Into<Shape<D>>>( shape: S, distribution: Distribution<B::FloatElem> ) -> Self
Create a random tensor of the given shape where each element is sampled from the given distribution.
sourcepub fn random_device<S: Into<Shape<D>>>(
shape: S,
distribution: Distribution<B::FloatElem>,
device: &B::Device
) -> Self
pub fn random_device<S: Into<Shape<D>>>( shape: S, distribution: Distribution<B::FloatElem>, 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.
sourcepub fn to_full_precision(&self) -> Tensor<B::FullPrecisionBackend, D>
pub fn to_full_precision(&self) -> Tensor<B::FullPrecisionBackend, D>
Returns a tensor with full precision based on the selected backend.
sourcepub fn from_full_precision(tensor: Tensor<B::FullPrecisionBackend, D>) -> Self
pub fn from_full_precision(tensor: Tensor<B::FullPrecisionBackend, D>) -> Self
Returns a tensor on the selected backend from a full precision tensor.
sourcepub fn detach(self) -> Self
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.
sourcepub fn require_grad(self) -> Self
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.
sourcepub fn is_require_grad(&self) -> bool
pub fn is_require_grad(&self) -> bool
Returns true if the tensor requires gradients during the backward pass.
sourcepub fn set_require_grad(self, require_grad: bool) -> Self
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.
source§impl<const D: usize, B: ADBackend> Tensor<B, D>
impl<const D: usize, B: ADBackend> Tensor<B, D>
sourcepub fn grad(&self, grads: &B::Gradients) -> Option<Tensor<B::InnerBackend, D>>
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.
sourcepub fn grad_remove(
&self,
grads: &mut B::Gradients
) -> Option<Tensor<B::InnerBackend, D>>
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.
sourcepub fn grad_replace(
&self,
grads: &mut B::Gradients,
grad: Tensor<B::InnerBackend, D>
)
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.
sourcepub fn inner(self) -> Tensor<B::InnerBackend, D>
pub fn inner(self) -> Tensor<B::InnerBackend, D>
Returns the inner tensor without the autodiff information.
sourcepub fn from_inner(inner: Tensor<B::InnerBackend, D>) -> Self
pub fn from_inner(inner: Tensor<B::InnerBackend, D>) -> Self
source§impl<B> Tensor<B, 1, Int>where
B: Backend,
impl<B> Tensor<B, 1, Int>where B: Backend,
sourcepub fn arange(range: Range<usize>) -> Self
pub fn arange(range: Range<usize>) -> Self
Returns a new integer tensor on the default device.
Arguments
range
- The range of values to generate.
sourcepub fn arange_step(range: Range<usize>, step: usize) -> Self
pub fn arange_step(range: Range<usize>, step: usize) -> Self
Returns a new integer tensor on the default device.
Arguments
range
- The range of values to generate.step
- The step between each value.
source§impl<const D: usize, B> Tensor<B, D, Int>where
B: Backend,
impl<const D: usize, B> Tensor<B, D, Int>where B: Backend,
sourcepub fn from_ints<A: Into<Data<i32, D>>>(ints: A) -> Self
pub fn from_ints<A: Into<Data<i32, D>>>(ints: A) -> Self
Create a tensor from integers (i32).
Example
use burn_tensor::backend::Backend;
use burn_tensor::{Tensor, Int};
fn example<B: Backend>() {
let _x: Tensor<B, 1, Int> = Tensor::from_ints([1, 2]);
let _y: Tensor<B, 2, Int> = Tensor::from_ints([[1, 2], [3, 4]]);
}
sourcepub fn float(self) -> Tensor<B, D, Float>
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 int_tensor = Tensor::<B, 1, Int>::arange(0..5);
let float_tensor = int_tensor.float();
}
source§impl<B, const D: usize, K> Tensor<B, D, K>where
B: Backend,
K: Numeric<B>,
K::Elem: Element,
impl<B, const D: usize, K> Tensor<B, D, K>where B: Backend, K: Numeric<B>, K::Elem: Element,
sourcepub fn into_scalar(self) -> K::Elem
pub fn into_scalar(self) -> K::Elem
sourcepub fn add_scalar<E: ElementConversion>(self, other: E) -> Self
pub fn add_scalar<E: ElementConversion>(self, other: E) -> Self
Applies element wise addition operation with a scalar.
y = x + s
sourcepub fn sub_scalar<E: ElementConversion>(self, other: E) -> Self
pub fn sub_scalar<E: ElementConversion>(self, other: E) -> Self
Applies element wise subtraction operation with a scalar.
y = x - s
sourcepub fn div_scalar<E: ElementConversion>(self, other: E) -> Self
pub fn div_scalar<E: ElementConversion>(self, other: E) -> Self
Applies element wise division operation with a scalar.
y = x / s
sourcepub fn mul(self, other: Self) -> Self
pub fn mul(self, other: Self) -> Self
Applies element wise multiplication operation.
y = x2 * x1
sourcepub fn mul_scalar<E: ElementConversion>(self, other: E) -> Self
pub fn mul_scalar<E: ElementConversion>(self, other: E) -> Self
Applies element wise multiplication operation with a scalar.
y = x * s
sourcepub fn zeros<S: Into<Shape<D>>>(shape: S) -> Self
pub fn zeros<S: Into<Shape<D>>>(shape: S) -> Self
Create a tensor of the given shape where each element is zero.
sourcepub fn zeros_device<S: Into<Shape<D>>>(shape: S, device: &B::Device) -> Self
pub fn zeros_device<S: Into<Shape<D>>>(shape: S, device: &B::Device) -> Self
Create a tensor of the given shape where each element is zero.
sourcepub fn ones<S: Into<Shape<D>>>(shape: S) -> Self
pub fn ones<S: Into<Shape<D>>>(shape: S) -> Self
Create a tensor of the given shape where each element is one.
sourcepub fn ones_device<S: Into<Shape<D>>>(shape: S, device: &B::Device) -> Self
pub fn ones_device<S: Into<Shape<D>>>(shape: S, device: &B::Device) -> Self
Create a tensor of the given shape where each element is one.
sourcepub fn full<S: Into<Shape<D>>, E: ElementConversion>(
shape: S,
fill_value: E
) -> Self
pub fn full<S: Into<Shape<D>>, E: ElementConversion>( shape: S, fill_value: E ) -> Self
Create a tensor of the given shape where each element is equal to the provided value.
sourcepub fn full_device<S: Into<Shape<D>>, E: ElementConversion>(
shape: S,
fill_value: E,
device: &B::Device
) -> Self
pub fn full_device<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.
sourcepub fn mean(self) -> Tensor<B, 1, K>
pub fn mean(self) -> Tensor<B, 1, K>
Aggregate all elements in the tensor with the mean operation.
sourcepub fn sum(self) -> Tensor<B, 1, K>
pub fn sum(self) -> Tensor<B, 1, K>
Aggregate all elements in the tensor with the sum operation.
sourcepub fn mean_dim(self, dim: usize) -> Self
pub fn mean_dim(self, dim: usize) -> Self
Aggregate all elements along the given dimension or axis in the tensor with the mean operation.
sourcepub fn sum_dim(self, dim: usize) -> Self
pub fn sum_dim(self, dim: usize) -> Self
Aggregate all elements along the given dimension or axis in the tensor with the sum operation.
sourcepub fn equal_elem<E: Element>(self, other: E) -> Tensor<B, D, Bool>
pub fn equal_elem<E: Element>(self, other: E) -> Tensor<B, D, Bool>
Applies element wise equal comparison and returns a boolean tensor.
sourcepub fn greater(self, other: Self) -> Tensor<B, D, Bool>
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.
sourcepub fn greater_equal(self, other: Self) -> Tensor<B, D, Bool>
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.
sourcepub fn lower(self, other: Self) -> Tensor<B, D, Bool>
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.
sourcepub fn lower_equal(self, other: Self) -> Tensor<B, D, Bool>
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.
sourcepub fn greater_elem<E: ElementConversion>(self, other: E) -> Tensor<B, D, Bool>
pub fn greater_elem<E: ElementConversion>(self, other: E) -> Tensor<B, D, Bool>
Applies element wise greater comparison and returns a boolean tensor.
sourcepub fn greater_equal_elem<E: ElementConversion>(
self,
other: E
) -> Tensor<B, D, Bool>
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.
sourcepub fn lower_elem<E: ElementConversion>(self, other: E) -> Tensor<B, D, Bool>
pub fn lower_elem<E: ElementConversion>(self, other: E) -> Tensor<B, D, Bool>
Applies element wise lower comparison and returns a boolean tensor.
sourcepub fn lower_equal_elem<E: ElementConversion>(
self,
other: E
) -> Tensor<B, D, Bool>
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.
sourcepub fn mask_where(self, mask: Tensor<B, D, Bool>, value: Self) -> Self
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.
sourcepub fn mask_fill<E: ElementConversion>(
self,
mask: Tensor<B, D, Bool>,
value: E
) -> Self
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.
sourcepub fn gather(self, dim: usize, indices: Tensor<B, D, Int>) -> Self
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.
sourcepub fn scatter(
self,
dim: usize,
indices: Tensor<B, D, Int>,
values: Self
) -> Self
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.
sourcepub fn select(self, dim: usize, indices: Tensor<B, 1, Int>) -> Self
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
sourcepub fn select_assign(
self,
dim: usize,
indices: Tensor<B, 1, Int>,
values: Tensor<B, D, K>
) -> Self
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
sourcepub fn argmax(self, dim: usize) -> Tensor<B, D, Int>
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 tensor = Tensor::<B, 3>::ones(Shape::new([2, 3, 3]));
let tensor = tensor.argmax(1);
println!("{:?}", tensor.shape());
// Shape { dims: [2, 1, 3] }
}
sourcepub fn max_dim(self, dim: usize) -> Tensor<B, D, K>
pub fn max_dim(self, dim: usize) -> Tensor<B, D, K>
Find the maximum value along the given dimension.
sourcepub fn max_dim_with_indices(
self,
dim: usize
) -> (Tensor<B, D, K>, Tensor<B, D, Int>)
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.
sourcepub fn argmin(self, dim: usize) -> Tensor<B, D, Int>
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 tensor = Tensor::<B, 3>::ones(Shape::new([2, 3, 3]));
let tensor = tensor.argmin(1);
println!("{:?}", tensor.shape());
// Shape { dims: [2, 1, 3] }
}
sourcepub fn min_dim(self, dim: usize) -> Tensor<B, D, K>
pub fn min_dim(self, dim: usize) -> Tensor<B, D, K>
Find the minimum value along the given dimension.
sourcepub fn min_dim_with_indices(
self,
dim: usize
) -> (Tensor<B, D, K>, Tensor<B, D, Int>)
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.
sourcepub fn clamp<E: ElementConversion>(self, min: E, max: E) -> Self
pub fn clamp<E: ElementConversion>(self, min: E, max: E) -> Self
sourcepub fn clamp_min<E: ElementConversion>(self, min: E) -> Self
pub fn clamp_min<E: ElementConversion>(self, min: E) -> Self
sourcepub fn clamp_max<E: ElementConversion>(self, max: E) -> Self
pub fn clamp_max<E: ElementConversion>(self, max: E) -> Self
Trait Implementations§
source§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,
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,
source§impl<B, const D: usize, K> Add for Tensor<B, D, K>where
B: Backend,
K: Numeric<B>,
K::Elem: Element,
impl<B, const D: usize, K> Add for Tensor<B, D, K>where B: Backend, K: Numeric<B>, K::Elem: Element,
source§impl<B, const D: usize, K> Clone for Tensor<B, D, K>where
B: Backend + Clone,
K: TensorKind<B> + Clone,
K::Primitive<D>: Clone,
impl<B, const D: usize, K> Clone for Tensor<B, D, K>where B: Backend + Clone, K: TensorKind<B> + Clone, K::Primitive<D>: Clone,
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,
impl<B, const D: usize, K> Debug for Tensor<B, D, K>where B: Backend + Debug, K: TensorKind<B> + Debug, K::Primitive<D>: Debug,
source§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,
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