17 releases (breaking)

0.16.0 Feb 8, 2025
0.15.0 Dec 28, 2024
0.14.1 Nov 16, 2024
0.12.0 Jul 30, 2024
0.1.0 Dec 31, 2023

#537 in Machine learning

Download history 487/week @ 2024-11-01 420/week @ 2024-11-08 622/week @ 2024-11-15 682/week @ 2024-11-22 544/week @ 2024-11-29 606/week @ 2024-12-06 578/week @ 2024-12-13 304/week @ 2024-12-20 433/week @ 2024-12-27 603/week @ 2025-01-03 548/week @ 2025-01-10 488/week @ 2025-01-17 461/week @ 2025-01-24 482/week @ 2025-01-31 658/week @ 2025-02-07 253/week @ 2025-02-14

1,934 downloads per month
Used in 10 crates (7 directly)

MIT/Apache

355KB
8K SLoC

rten-tensor

rten-tensor is the foundational library that provides multi-dimensional arrays used by RTen. It is similar to ndarray but tailored for use in the RTen library.


lib.rs:

rten_tensor provides multi-dimensional arrays, commonly referred to as tensors in a machine learning context.

Storage and layout

A tensor is a combination of data storage and a layout. The data storage determines the element type and how the data is owned. A tensor can be:

  • Owned (like Vec<T>)
  • Borrowed (like &[T] or &mut [T])
  • Maybe-owned (like Cow[T])

The layout determines the number of dimensions (the rank) and size of each (the shape) and how indices map to offsets in the storage. The dimension count can be static (fixed at compile time) or dynamic (variable at runtime).

Tensor types and traits

The base type for all tensors is [TensorBase]. This is not normally used directly but instead via a type alias which specifies the data ownership and layout:

Rank Owned Borrowed Mutably borrowed Owned or borrowed
Static [NdTensor] [NdTensorView] [NdTensorViewMut] [CowNdTensor]
Dynamic [Tensor] [TensorView] [TensorViewMut] [CowTensor]

All tensors implement the [Layout] trait, which provide methods to query the shape, dimension count and strides of the tensor. Tensor views provide various methods for indexing, iterating, slicing and transforming them. The [AsView] trait provides access to these methods for owned and mutably borrowed tensors. Conceptually it is similar to how Deref allows accesing methods for &[T] on a Vec<T>. The preferred way to import the traits is via the prelude:

use rten_tensor::prelude::*;
use rten_tensor::NdTensor;

let tensor = NdTensor::from([[1, 2], [3, 4]]);

let transposed_elems: Vec<_> = tensor.transposed().iter().copied().collect();
assert_eq!(transposed_elems, [1, 3, 2, 4]);

Serialization

Tensors can be serialized and deserialized using serde if the serde feature is enabled. The serialized representation of a tensor includes its shape and elements in row-major (C) order. The JSON serialization of a matrix (NdTensor<f32, 2>) looks like this for example:

{
  "shape": [2, 2],
  "data": [0.5, 1.0, 1.5, 2.0]
}

Dependencies

~220KB