flashinfer.norm.layernorm

flashinfer.norm.layernorm(input: Tensor, gemma: Tensor, beta: Tensor, eps: float = 1e-06) → Tensor

层归一化。 :param input: 输入张量，形状为 (batch_size, hidden_size)。需要为 bfloat16 类型。 :type input: torch.Tensor :param gemma: Gemma 张量，形状为 (hidden_size,)。需要为 float32 类型。 :type gemma: torch.Tensor :param beta: Beta 张量，形状为 (hidden_size,)。需要为 float32 类型。 :type beta: torch.Tensor :param eps: 用于数值稳定的 epsilon 值。 :type eps: float

返回值:

output – 层归一化后的张量，形状为 (batch_size, hidden_size)。与输入张量具有相同的 dtype。

返回值类型:

torch.Tensor