Struct StridedTensor

Source

pub struct StridedTensor<T: Clone + Default + Copy + 'static> {
    pub underlying_data: Array2<T>,
    pub inner_dim: usize,
    pub element_lengths: Array1<i64>,
    pub max_element_len: usize,
    pub precomputed_strides: Vec<usize>,
    pub cumulative_lengths: Array1<i64>,
}

Expand description

A data structure for efficient batch lookups on tensors of varying lengths.

StridedTensor stores a collection of variable-length sequences as a single, contiguous array with padding. It precomputes several views with different strides to optimize retrieval of sequences with common lengths.

Fields§

§underlying_data: Array2<T>

The flattened, contiguous data containing all sequences with padding

§inner_dim: usize

The shape of each individual element within the data (e.g., embedding dim)

§element_lengths: Array1<i64>

Length of each element sequence

§max_element_len: usize

Maximum length found among all sequences

§precomputed_strides: Vec<usize>

Sorted vector of strides for precomputed views

§cumulative_lengths: Array1<i64>

Cumulative sum of element_lengths for offset calculation

Implementations§

Source §

impl<T: Clone + Default + Copy + 'static> StridedTensor<T>

Source

pub fn new(data: Array2<T>, lengths: Array1<i64>) -> Self

Creates a new StridedTensor from concatenated data and lengths.

§Arguments

data - Concatenated data of all elements, shape [total_tokens, inner_dim]
lengths - Length of each element sequence

§Returns

A new StridedTensor with precomputed views for efficient lookup

Source

pub fn len(&self) -> usize

Returns the number of elements (sequences) stored

Source

pub fn is_empty(&self) -> bool

Returns true if there are no elements

Source

pub fn total_tokens(&self) -> usize

Get the total number of tokens across all sequences

Source §

impl StridedTensor<i64>

Source

pub fn lookup_1d(&self, indices: &[usize]) -> (Array1<i64>, Array1<i64>)

Retrieve elements by their indices (for codes - 1D per element).

§Arguments

indices - Indices of elements to retrieve

§Returns

Tuple of (flattened data, lengths for each element)

Source §

impl StridedTensor<u8>

Source

pub fn lookup_2d(&self, indices: &[usize]) -> (Array2<u8>, Array1<i64>)

Retrieve elements by their indices (for residuals - 2D per element).

§Arguments

indices - Indices of elements to retrieve

§Returns

Tuple of (concatenated data, lengths for each element)

Source §

impl StridedTensor<usize>

Source

pub fn lookup_codes(&self, indices: &[usize]) -> (Array1<usize>, Array1<i64>)

Retrieve elements by their indices (for codes stored as usize).

§Arguments

indices - Indices of elements to retrieve

§Returns

Tuple of (flattened codes, lengths for each element)

Trait Implementations§

Source §

impl<T: Clone + Clone + Default + Copy + 'static> Clone for StridedTensor<T>

Source §

fn clone(&self) -> StridedTensor<T>

Returns a duplicate of the value. Read more

1.0.0 · Source§

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

Performs copy-assignment from source. Read more

Auto Trait Implementations§

§

impl<T> UnwindSafe for StridedTensor<T>
where T: RefUnwindSafe,

Blanket Implementations§

Source §

impl<T> Any for T
where T: 'static + ?Sized,

Source §

fn type_id(&self) -> TypeId

Gets the TypeId of self. Read more

Source §

impl<T> Borrow<T> for T
where T: ?Sized,

Source §

fn borrow(&self) -> &T

Immutably borrows from an owned value. Read more

Source §

impl<T> BorrowMut<T> for T
where T: ?Sized,

Source §

fn borrow_mut(&mut self) -> &mut T

Mutably borrows from an owned value. Read more

Source §

impl<T> CloneToUninit for T
where T: Clone,

Source §

unsafe fn clone_to_uninit(&self, dest: *mut u8)

🔬This is a nightly-only experimental API. (clone_to_uninit)

Performs copy-assignment from self to dest. Read more

Source §

impl<T> From<T> for T

Source §

fn from(t: T) -> T

Returns the argument unchanged.

Source §

impl<T, U> Into for T
where U: From<T>,

Source §

fn into(self) -> U

Calls U::from(self).

That is, this conversion is whatever the implementation of From<T> for U chooses to do.

Source §

impl<T> IntoEither for T

Source §

fn into_either(self, into_left: bool) -> Either<Self, Self>

Converts self into a Left variant of Either<Self, Self> if into_left is true. Converts self into a Right variant of Either<Self, Self> otherwise. Read more

Source §

fn into_either_with<F>(self, into_left: F) -> Either<Self, Self>
where F: FnOnce(&Self) -> bool,

Converts self into a Left variant of Either<Self, Self> if into_left(&self) returns true. Converts self into a Right variant of Either<Self, Self> otherwise. Read more

Source §