LayerNorm 前向实现：数值稳定性陷阱

import numpy as np

class LayerNorm:
    def __init__(self, normalized_shape, eps=1e-5):
        self.normalized_shape = tuple(
            normalized_shape if isinstance(normalized_shape, (tuple, list))
            else (normalized_shape,)
        )
        self.eps = eps
        self.gamma = np.ones(self.normalized_shape, dtype=np.float32)
        self.beta = np.zeros(self.normalized_shape, dtype=np.float32)

    def forward(self, x):
        if tuple(x.shape[-len(self.normalized_shape):]) != self.normalized_shape:
            raise ValueError("input tail shape does not match normalized_shape")
        axes = tuple(range(x.ndim - len(self.normalized_shape), x.ndim))

        input_dtype = x.dtype
        input_is_bfloat16 = input_dtype.name == "bfloat16"
        input_is_float = input_is_bfloat16 or np.issubdtype(
            input_dtype, np.floating
        )
        low_precision = input_dtype == np.float16 or input_is_bfloat16
        if low_precision or not input_is_float:
            work = x.astype(np.float32)
        else:
            work = x

        mean = np.mean(work, axis=axes, keepdims=True)
        var = np.mean((work - mean) ** 2, axis=axes, keepdims=True)
        x_hat = (work - mean) / np.sqrt(var + self.eps)
        gamma = self.gamma.astype(work.dtype, copy=False)
        beta = self.beta.astype(work.dtype, copy=False)
        out = x_hat * gamma + beta
        return out.astype(input_dtype, copy=False) if input_is_float else out