Skip to main content

TensorBatchOps

Trait TensorBatchOps 

Source
pub trait TensorBatchOps: Send + Sync {
    // Required methods
    fn batch_matmul(
        &self,
        a_batch: &[&TensorRef],
        b_batch: &[&TensorRef],
    ) -> Result<Vec<TensorRef>>;
    fn stack(&self, tensors: &[&TensorRef], dim: usize) -> Result<TensorRef>;
    fn unstack(&self, tensor: &TensorRef, dim: usize) -> Result<Vec<TensorRef>>;
    fn pad_batch(
        &self,
        tensors: &[&TensorRef],
        target_shape: &[usize],
    ) -> Result<Vec<TensorRef>>;
}
Expand description

Tensor batch operations for efficient processing

Required Methods§

Source

fn batch_matmul( &self, a_batch: &[&TensorRef], b_batch: &[&TensorRef], ) -> Result<Vec<TensorRef>>

Batch matrix multiplication for multiple pairs

Source

fn stack(&self, tensors: &[&TensorRef], dim: usize) -> Result<TensorRef>

Stack tensors along new dimension

Source

fn unstack(&self, tensor: &TensorRef, dim: usize) -> Result<Vec<TensorRef>>

Unstack tensor along specified dimension

Source

fn pad_batch( &self, tensors: &[&TensorRef], target_shape: &[usize], ) -> Result<Vec<TensorRef>>

Pad tensors in batch to same shape

Implementors§