Mojo function

layer_norm_cpu

layer_norm_cpu[dtype: DType, //, input_fn: fn[Int](Int, Int) capturing -> SIMD[dtype, $0], gamma_fn: fn[Int, Int](IndexList[$1]) capturing -> SIMD[dtype, $0], output_fn: fn[Int, Int](row: Int, col: Int, val: SIMD[dtype, $0]) capturing -> None](num_rows: Int, num_cols: Int, beta: NDBuffer[dtype, 1, origin], epsilon: SIMD[dtype, 1])

Computes layernorm(elementwise_fn(x)) across the last dimension of x, where layernorm is defined as $(x-mean(x))/(sqrt(var(x)+eps)*gamma_fn + beta$ .

Currently performs 3 passes over the input data. This can be reduced to 2 by fusing the add, mean, and variance loops using Welford's algorithm.

Parameters:

dtype (DType): The x and out buffers' elements dtype.
input_fn (fn[Int](Int, Int) capturing -> SIMD[dtype, $0]): Function called to generate an input value.
gamma_fn (fn[Int, Int](IndexList[$1]) capturing -> SIMD[dtype, $0]): Function called to generate a gamma value.
output_fn (fn[Int, Int](row: Int, col: Int, val: SIMD[dtype, $0]) capturing -> None): Function called to store the output value.

Args:

num_rows (Int): The number of rows in the input tensor.
num_cols (Int): The number of columns in the input tensor.
beta (NDBuffer[dtype, 1, origin]): The beta value to use in the layernorm calculation.
epsilon (SIMD[dtype, 1]): The eps value to use in the layernorm calculation.

layer_norm_cpu[dtype: DType, rank: Int, //, input_fn: fn[Int, Int](IndexList[$1]) capturing -> SIMD[dtype, $0], gamma_fn: fn[Int, Int](IndexList[$1]) capturing -> SIMD[dtype, $0], output_fn: fn[Int, Int, Int](idx: IndexList[$1], val: SIMD[dtype, $0]) capturing -> None](shape: IndexList[rank], beta: NDBuffer[dtype, 1, origin], epsilon: SIMD[dtype, 1])