Struct TensorOps 

pub struct TensorOps { /* private fields */ }

Bounded shape-specialized device tensor operations in one shared context.

Implementations

impl TensorOps

pub fn linear( &self, input: &DeviceTensor, weight: &DeviceTensor, output: DType, ) -> Result<DeviceTensor>

Single-row BF16 x * weightᵀ, accumulating in F32. Weights are contiguous [N, K]; input is [1, K] with K a multiple of 128. Output may be BF16 or F32. This bandwidth-oriented decode operation does not accept biases.

pub fn linear_many( &self, input: &DeviceTensor, weights: &[DeviceTensor], output: DType, ) -> Result<Vec<DeviceTensor>>

One dispatch for one, two or three single-row projections, e.g. Q/K/V or gate/up. Weights and output widths can differ; K and output dtype agree. Returned tensors retain a shared plan-slot lease. Three live results per shape are supported; release one before requesting a fourth.

pub fn linear_fragment( &self, input: &TensorDesc, weights: &[TensorDesc], output: DType, ) -> Result<ModelFragment>

Record the same decode projections in a caller-owned inference graph. Fragment inputs are [input, weights...], outputs follow weight order.

impl TensorOps

pub fn gather_rows_fragment( &self, input: &TensorDesc, rows: usize, ) -> Result<ModelFragment>

Recordable row gather with directly bound source and U32 [rows] indices. Output is host-visible so final records can be mapped without a copy; GPU consumers can bind it directly. Out-of-range indices select row zero, so callers must validate indices before publication.

pub fn new(context: &ModelContext, capacity: usize) -> Result<Self>

Create lazy bounded caches, evicting only unleased and drained plans.

pub fn gather_rows( &self, input: &DeviceTensor, indices: &[usize], ) -> Result<DeviceTensor>

Gather complete contiguous rows using checked host indices. Duplicates and arbitrary ordering are allowed. Only indices are uploaded; row bytes stay on device. Three private slots per power-of-two output-size bucket bound preparation churn; the returned view retains its slot until dropped.

pub fn prepare_finite(&self, input: &TensorDesc) -> Result<Arc<PreparedModel>>

Reduce contiguous F32 to one I32 error flag: zero iff every value is finite. Empty tensors are not accepted. Multi-pass bounded reductions avoid atomics and leave input/output ownership with normal slot leases.

pub fn finite_fragment(&self, input: &TensorDesc) -> Result<ModelFragment>

Record a finite-value reduction directly into a caller-owned graph. The output is one I32 flag with the same contract as prepare_finite.

pub fn finite(&self, input: &DeviceTensor) -> Result<Inference>

Return a resident non-finite flag without downloading the input tensor.