LayerNorm

import torch
import torch.nn as nn

#使能KDNN
torch._C._set_kdnn_enabled(True)

# 输入数据: (batch_size, seq_len, hidden_dim)
input = torch.randn(2, 5, 10)  # 2个样本，5个时间步，特征维度10 默认为fp32数据类型

# 定义 LayerNorm (对最后一维 hidden_dim=10 归一化)
layer_norm = nn.LayerNorm(10)  # 或者 normalized_shape=[5,10] 对最后两维归一化 默认为fp32数据类型

# 前向计算
output = layer_norm(input)

# 验证: 归一化后最后一维的均值接近0，方差接近1
print("LayerNorm 输出均值:", output.mean(dim=-1))  # 应接近0
print("LayerNorm 输出方差:", output.var(dim=-1, unbiased=False))  # 应接近1