For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo function
composite_rms_norm_rope_shape
def composite_rms_norm_rope_shape[dtype: DType](input: T, weight: T, epsilon: Float32, weight_offset: Scalar[dtype], cos_vals: T, sin_vals: T) -> IndexList[T.rank]
Computes the output shape for the mo.composite.rms_norm_rope graph op.
Args:
- input (
T): Activation tensor normalized by RMS then rotated by RoPE; the last dimension must be even. - weight (
T): Per-column scale weights applied after RMS normalization. - epsilon (
Float32): Small constant added inside the RMS normalization square root for numerical stability. - weight_offset (
Scalar[dtype]): Scalar offset added toweightbefore scaling. - cos_vals (
T): Cosine table used by the RoPE rotation, matchinginputrank. - sin_vals (
T): Sine table used by the RoPE rotation, matchinginputrank.
Returns:
IndexList[T.rank]: The output shape, which matches the input shape.