Mojo function

layer_norm_shape

layer_norm_shape[type: DType, rank: Int, single_thread_blocking_override: Bool](input: NDBuffer[type, rank, origin], gamma: NDBuffer[type, 1, origin, __init__[::Intable](1)], beta: NDBuffer[type, 1, origin, __init__[::Intable](1)], epsilon: SIMD[type, 1]) -> Index[rank]

Compute the output shape of a layer_norm operation.

Parameters:

type (DType): Type of the input tensors.
rank (Int): Rank of the input tensor.
single_thread_blocking_override (Bool): If True, then the operation is run synchronously using a single thread.

Args:

input (NDBuffer[type, rank, origin]): The input tensor.
gamma (NDBuffer[type, 1, origin, __init__[::Intable](1)]): The tensor for gamma coefficient.
beta (NDBuffer[type, 1, origin, __init__[::Intable](1)]): The tensor for beta coefficient.
epsilon (SIMD[type, 1]): The tensor for epsilon coefficient.

Returns:

The output shape.