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

rms_norm_rope

def rms_norm_rope[input_dtype: DType, output_dtype: DType, cos_sin_dtype: DType, rank: Int, InputFn: def[width: Int, alignment: Int, coord_rank: Int](IndexList[coord_rank]) -> SIMD[input_dtype, width] & RegisterPassable & ImplicitlyCopyable, CosFn: def[width: Int, alignment: Int, coord_rank: Int](IndexList[coord_rank]) -> SIMD[cos_sin_dtype, width] & RegisterPassable & ImplicitlyCopyable, SinFn: def[width: Int, alignment: Int, coord_rank: Int](IndexList[coord_rank]) -> SIMD[cos_sin_dtype, width] & RegisterPassable & ImplicitlyCopyable, OutputFn: def[width: SIMDLength, rank: Int, alignment: Int](IndexList[rank], SIMD[output_dtype, width]) -> None & RegisterPassable & ImplicitlyCopyable, AxisSizeT: CoordLike, /, target: StringSpan[ImmStaticOrigin], multiply_before_cast: Bool = True, reduce_dim: Int = (rank - Int(1))](input_fn: InputFn, cos_fn: CosFn, sin_fn: SinFn, output_fn: OutputFn, shape: Coord, axis_size: AxisSizeT, gamma: TileTensor[input_dtype, Storage=gamma.Storage, address_space=gamma.address_space, linear_idx_type=gamma.linear_idx_type], epsilon: Scalar[input_dtype], weight_offset: Scalar[input_dtype], context: Optional[DeviceContext] = None) where (eq InputFn.input_dtype, input_dtype) where (eq CosFn.cos_sin_dtype, cos_sin_dtype) where (eq SinFn.cos_sin_dtype, cos_sin_dtype) where (eq OutputFn.output_dtype, output_dtype)