rms_norm(x, weight, eps, *[, stream])
|
Root Mean Square normalization (RMS norm). |
layer_norm(x, weight, bias, eps, *[, stream])
|
Layer normalization. |
rope(a, dims, *, traditional, base, scale, ...)
|
Apply rotary positional encoding to the input. |
scaled_dot_product_attention(q, k, v, *, scale)
|
A fast implementation of multi-head attention: O = softmax(Q @ K.T, dim=-1) @ V. |
affine_quantize(w, /, scales, biases[, ...])
|
Quantize the matrix w using the provided scales and biases and the group_size and bits configuration. |
metal_kernel(name, input_names, ...[, ...])
|
A jit-compiled custom Metal kernel defined from a source string. |
