Expand description
Convenience functions for tensor operations
Functions§
- add
- attention
- causal_
mask - Create a causal attention mask (lower triangular)
- causal_
sliding_ window_ mask - Create a combined causal + sliding window mask This is the typical mask used in Mistral/Mixtral
- clamp
- Clamp values to a range
- concat
- conv1d
- 1D convolution
- embedding
- exp
- flash_
attention - Fused scaled dot-product attention (Flash Attention pattern) Computes: softmax(Q @ K^T / sqrt(d_k)) @ V Works for: LLaMA, Qwen, Gemma, Mistral, Phi, and all attention-based models
- fused_
residual_ rms_ norm - Fused residual add + RMS normalization Computes: rms_norm(residual + hidden, weight, eps) Used by: All transformer models with pre-normalization
- fused_
swiglu - Fused SwiGLU activation: silu(gate) * up Used by: LLaMA, Qwen, Mistral, and other modern transformer MLPs
- gather
- Gather elements along dimension
- gelu
- layer_
norm - matmul
- mul
- normalize
- randn
- rms_
norm - scale
- scatter
- Scatter elements along dimension
- sigmoid
- Sigmoid activation
- silu
- sliding_
window_ mask - Create a sliding window attention mask Returns a mask where 1.0 means “attend” and 0.0 means “don’t attend” Each position can only attend to positions within window_size positions before it
- softmax
- Softmax along specified dimension
- sub
- Element-wise subtraction
- tanh
- Tanh activation
- topk
- Top-k operation
- transpose
- Transpose tensor (swap last two dimensions)
- transpose_
dims - Transpose tensor with specific dimensions
- zeros