xma.functional.rmsnorm

rmsnorm(x: Tensor, weight: Tensor | None, eps: float | None = None, memory_efficient: bool = False, deterministic: bool = False, *, kernel_backend: KernelBackend | None = None) → Tensor

RMSNorm computation

Parameters:

• x (torch.Tensor) – input activation

• weight (torch.Tensor | None) – RMSNorm weight

• eps (float | None) – epsilon. Defaults to None.

• memory_efficient (bool) – memory efficient = False caches RMSNorm’s denominator in the forward. Defaults to False.

• deterministic (bool) – whether to use deterministic backward. Defaults to False.

• kernel_backend (KernelBackend | None) – KernelBackend

Returns:

output tensor

Return type:

Tensor