Mojo function

layer_norm_cpu

layer_norm_cpu[type: DType, //, input_fn: fn[Int](Int, Int) capturing -> SIMD[type, $0], gamma_fn: fn[Int, Int](Index[$1]) capturing -> SIMD[type, $0]](out_buf: NDBuffer[type, 2, origin, shape], beta: NDBuffer[type, 1, origin], epsilon: SIMD[type, 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:

type (DType): The x and out buffers' elements dtype.
input_fn (fn[Int](Int, Int) capturing -> SIMD[type, $0]): Function called to generate an input value.
gamma_fn (fn[Int, Int](Index[$1]) capturing -> SIMD[type, $0]): Function called to generate a gamma value.

Args:

out_buf (NDBuffer[type, 2, origin, shape]): The output buffer.
beta (NDBuffer[type, 1, origin]): The beta value to use in the layernorm calculation.
epsilon (SIMD[type, 1]): The eps value to use in the layernorm calculation.

layer_norm_cpu[type: DType, rank: Int, //, input_fn: fn[Int, Int](Index[$1]) capturing -> SIMD[type, $0], gamma_fn: fn[Int, Int](Index[$1]) capturing -> SIMD[type, $0]](shape: Index[rank, element_type=element_type], beta: NDBuffer[type, 1, origin], epsilon: SIMD[type, 1], output: NDBuffer[type, rank, origin, shape, strides])