fix: more efficient implementation of layernorm (ivy-llc#23620)

iababio · Sep 27, 2023 · a830408 · a830408
1 parent 42ebaf8
commit a830408
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/ivy/functional/ivy/norms.py b/ivy/functional/ivy/norms.py
@@ -126,9 +126,8 @@ def layer_norm(
     """
     mean = ivy.mean(x, axis=normalized_idxs, keepdims=True)
     var = ivy.var(x, axis=normalized_idxs, keepdims=True)
-    x = ivy.divide(
-        ivy.add(ivy.negative(mean), x), ivy.stable_pow(var, 0.5, min_base=eps)
-    )
+
+    x = (x - mean) / (var + eps) ** 0.5
 
     if scale is not None:
         if offset is not None: