compute_mean_rstd() — pytorch Function Reference

Architecture documentation for the compute_mean_rstd() function in kernels.py from the pytorch codebase.

Function python DynamoBenchmarks Microbenchmarking called by 1

Entity Profile

DynamoBenchmarks→ Microbenchmarking→ compute_mean_rstd() — pytorch Function Reference

Dependency Diagram

graph TD
  94e8675f_d7e9_407b_4e04_6f5632247962["compute_mean_rstd()"]
  ac72338f_8a94_8d37_87cb_d74c1927bca1["liger()"]
  ac72338f_8a94_8d37_87cb_d74c1927bca1 -->|calls| 94e8675f_d7e9_407b_4e04_6f5632247962
  style 94e8675f_d7e9_407b_4e04_6f5632247962 fill:#6366f1,stroke:#818cf8,color:#fff

Relationship Graph

Source Code

benchmarks/dynamo/genai_layers/kernels.py lines 701–706

    def compute_mean_rstd(self, x, eps):
        x = x.float()

        var, mean = torch.var_mean(x, dim=-1, keepdim=True, correction=0)
        rstd = torch.rsqrt(var + eps)
        return mean, rstd