layernorm is all axis but the first
MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -313,8 +313,9 @@ 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):
- y = (self - self.mean(axis=-1, keepdim=True))
- return y.div((y*y).mean(axis=-1, keepdim=True).add(eps).sqrt())
+ axis = range(1, len(self.shape))
+ y = (self - self.mean(axis=axis, keepdim=True))
+ return y.div((y*y).mean(axis=axis, keepdim=True).add(eps).sqrt())
# An instantiation of the Function is the Context
class Function:

layernorm is all axis but the first