From dec5334da9b8a34c8dfb00fe01b4ac2c774f9645 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Mon, 26 Sep 2022 10:11:38 -0400 Subject: [PATCH] revert layernorm to have axis param --- tinygrad/tensor.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 87555c4ef4..82e6ecf4f3 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -312,8 +312,7 @@ class Tensor: def sequential(self, ll:List[Callable[[Tensor], Tensor]]): return functools.reduce(lambda x,f: f(x), ll, self) - def layernorm(self, eps=1e-5): - axis = range(1, len(self.shape)) + def layernorm(self, axis=-1, eps=1e-5): y = (self - self.mean(axis=axis, keepdim=True)) return y.div((y*y).mean(axis=axis, keepdim=True).add(eps).sqrt())