one liner that's more clear
MODIFY
examples/stable_diffusion.py
examples/stable_diffusion.py
@@ -7,6 +7,7 @@ from extra.utils import fake_torch_load_zipped, get_child
from tinygrad.nn import Conv2d
from tinygrad.tensor import Tensor
+# TODO: rename to GroupNorm and put in nn.py
class Normalize:
def __init__(self, in_channels, num_groups=32):
self.weight = Tensor.uniform(in_channels)
@@ -14,19 +15,11 @@ class Normalize:
self.num_groups = num_groups
def __call__(self, x):
- x = x.reshape(x.shape[0], self.num_groups, x.shape[1]//self.num_groups, x.shape[2], x.shape[3])
+ # reshape for layernorm to work as group norm
+ # subtract mean and divide stddev
+ x = x.reshape(x.shape[0], self.num_groups, -1).layernorm().reshape(x.shape)
- # subtract mean
- x = x - x.mean(axis=(2,3,4), keepdim=True)
-
- # divide stddev
- eps = 1e-5
- x = x.div((x*x).mean(axis=(2,3,4), keepdim=True).add(eps).sqrt())
-
- # return to old shape
- x = x.reshape(x.shape[0], x.shape[1]*x.shape[2], x.shape[3], x.shape[4])
-
- # elementwise_affine
+ # elementwise_affine on channels
return (x * self.weight.reshape(1, -1, 1, 1)) + self.bias.reshape(1, -1, 1, 1)

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -260,11 +260,10 @@ class Tensor:
# ***** functional nn ops *****
- # TODO: fix the kwargs problem, then remove these
- # NOTE: perhaps don't, since they create NOOPs if the shape already matches. (now lazy should do this)
- def reshape(self, shape, *args): shape = argfix(shape, *args); return self._reshape(shape=shape) if tuple(self.shape) != shape else self
- def expand(self, shape, *args): shape = argfix(shape, *args); return self._expand(shape=shape) if tuple(self.shape) != shape else self
- def permute(self, order, *args): order = argfix(order, *args); return self._permute(order=order)
+ # TODO: fix the kwargs problem, then remove these (or not, since they now fix tuples)
+ def reshape(self, shape, *args): return self._reshape(shape=argfix(shape, *args))
+ def expand(self, shape, *args): return self._expand(shape=argfix(shape, *args))
+ def permute(self, order, *args): return self._permute(order=argfix(order, *args))
def linear(self, weight:Tensor, bias:Tensor):
shp = [1] * (len(self.shape)-1) + [-1]

one liner that's more clear