Struct web_rwkv::tensor::ops::TensorOp

pub struct TensorOp<'a> {
    pub pipeline: &'a ComputePipeline,
    pub bindings: Vec<BindGroup>,
    pub dispatch: [u32; 3],
}

Fields

pipeline: &'a ComputePipeline

bindings: Vec<BindGroup>

dispatch: [u32; 3]

Implementations

impl<'a> TensorOp<'a>

pub fn softmax(x: &'a TensorGpu<f32, ReadWrite>) -> Result<Self, TensorError>

Softmax operator applied on x.

pub fn layer_norm( w: &'a TensorGpu<f16, ReadWrite>, b: &'a TensorGpu<f16, ReadWrite>, x: &'a TensorGpu<f32, ReadWrite> ) -> Result<Self, TensorError>

Layer normalization applied on x, with weight w and bias b.

• x shape: [C, T, B].
• w shape: [C, 1, 1].
• b shape: [C, 1, 1].

pub fn group_norm( w: &'a TensorGpu<f16, ReadWrite>, b: &'a TensorGpu<f16, ReadWrite>, x: &'a TensorGpu<f32, ReadWrite> ) -> Result<Self, TensorError>

Group normalization applied on x, with weight w and bias b.

• x shape: [S, H, A].
• w shape: [S, H, 1].
• b shape: [S, H, 1].

pub fn matmul_vec_fp16( matrix: &'a TensorGpu<f16, ReadWrite>, input: TensorView<'a, f32>, output: TensorView<'a, f32> ) -> Result<Self, TensorError>

Fp32 matrix-vector multiplication.

• matrix shape: [C, R, 1].
• input shape: [C, T, B].
• output shape: [R, T, B].

pub fn matmul_vec_int8( matrix: &'a TensorGpu<u8, ReadWrite>, mx: &'a TensorGpu<f32, ReadWrite>, rx: &'a TensorGpu<f32, ReadWrite>, my: &'a TensorGpu<f32, ReadWrite>, ry: &'a TensorGpu<f32, ReadWrite>, input: TensorView<'a, f32>, output: TensorView<'a, f32> ) -> Result<Self, TensorError>

Int8 matrix-vector multiplication.

• matrix shape: [C, R, 1].
• mx and rx shape: [C, 1, 1].
• my and ry shape: [R, 1, 1].
• input shape: [C, T, B].
• output shape: [R, T, B].

pub fn matmul_mat_fp16( matrix: TensorView<'a, f16>, input: TensorView<'a, f16>, output: TensorView<'a, f32> ) -> Result<Self, TensorError>

Fp16 matrix-matrix multiplication.

• matrix shape: [K, M, B].
• input shape: [K, N, B].
• output shape: [M, N, B].

pub fn add( input: &'a TensorGpu<f32, ReadWrite>, output: &'a TensorGpu<f32, ReadWrite> ) -> Result<Self, TensorError>

Add input onto output.

pub fn token_shift( cursors: &'a TensorGpu<u32, ReadWrite>, time_mix: &'a TensorGpu<f16, ReadWrite>, x: &'a TensorGpu<f32, ReadWrite>, sx: TensorView<'_, f32>, output: &'a TensorGpu<f32, ReadWrite> ) -> Result<Self, TensorError>

pub fn time_mix( stack: &'a TensorGpu<u32, ReadWrite>, time_decay: &'a TensorGpu<f32, ReadWrite>, time_first: &'a TensorGpu<f32, ReadWrite>, k: &'a TensorGpu<f32, ReadWrite>, v: &'a TensorGpu<f32, ReadWrite>, r: &'a TensorGpu<f32, ReadWrite>, x: &'a TensorGpu<f32, ReadWrite>, state: TensorView<'_, f32> ) -> Result<Self, TensorError>

pub fn time_mix_v5( stack: &'a TensorGpu<u32, ReadWrite>, time_decay: &'a TensorGpu<f32, ReadWrite>, time_first: &'a TensorGpu<f32, ReadWrite>, k: &'a TensorGpu<f32, ReadWrite>, v: &'a TensorGpu<f32, ReadWrite>, r: &'a TensorGpu<f32, ReadWrite>, x: &'a TensorGpu<f32, ReadWrite>, state: TensorView<'_, f32> ) -> Result<Self, TensorError>

pub fn silu( input: &'a TensorGpu<f32, ReadWrite>, output: &'a TensorGpu<f32, ReadWrite> ) -> Result<Self, TensorError>

pub fn squared_relu( x: &'a TensorGpu<f32, ReadWrite> ) -> Result<Self, TensorError>

pub fn channel_mix( cursors: &'a TensorGpu<u32, ReadWrite>, r: &'a TensorGpu<f32, ReadWrite>, v: &'a TensorGpu<f32, ReadWrite>, x: &'a TensorGpu<f32, ReadWrite>, state: TensorView<'a, f32> ) -> Result<Self, TensorError>

pub fn blit( input: TensorView<'a, f32>, output: TensorView<'a, f32> ) -> Result<Self, TensorError>

Copy the content of input into output, given an offset.

pub fn blend( factor: &'a TensorGpu<f32, Uniform>, input: &'a TensorGpu<f32, ReadWrite>, output: &'a TensorGpu<f32, ReadWrite> ) -> Result<Self, TensorError>

pub fn quantize_mat_int8( input: &'a TensorGpu<f16, ReadWrite>, mx: &'a TensorGpu<f32, ReadWrite>, rx: &'a TensorGpu<f32, ReadWrite>, my: &'a TensorGpu<f32, ReadWrite>, ry: &'a TensorGpu<f32, ReadWrite>, output: &'a TensorGpu<u8, ReadWrite> ) -> Result<Vec<Self>, TensorError>

pub fn quantize_fp16( input: &'a TensorGpu<f32, ReadWrite>, output: &'a TensorGpu<f16, ReadWrite> ) -> Result<Self, TensorError>