add Linear to tinygrad.nn
MODIFY
examples/stable_diffusion.py
examples/stable_diffusion.py
@@ -8,7 +8,7 @@ import traceback
from tqdm import tqdm
from collections import namedtuple
from extra.utils import fake_torch_load_zipped, get_child
-from tinygrad.nn import Conv2d
+from tinygrad.nn import Conv2d, Linear
from tinygrad.tensor import Tensor
# TODO: refactor AttnBlock, CrossAttention, CLIPAttention to share code
@@ -35,7 +35,6 @@ class Normalize:
else:
return x.linear(self.weight, self.bias)
-
class AttnBlock:
def __init__(self, in_channels):
self.norm = Normalize(in_channels)
@@ -166,14 +165,6 @@ class AutoencoderKL:
latent = self.post_quant_conv(latent)
return self.decoder(latent)
-class Linear:
- def __init__(self, in_features, out_features, bias=True):
- self.weight = Tensor.empty(out_features, in_features)
- self.bias = Tensor.empty(out_features) if bias else None
-
- def __call__(self, x):
- return x.linear(self.weight.transpose(), self.bias)
-
# not to be confused with ResnetBlock
class ResBlock:
def __init__(self, channels, emb_channels, out_channels):

MODIFY
test/test_nn.py
test/test_nn.py
@@ -57,14 +57,14 @@ class TestNN(unittest.TestCase):
def _test_linear(x):
# create in tinygrad
- layer = (Tensor.uniform(in_dim, out_dim), Tensor.zeros(out_dim))
- z = x.linear(*layer)
+ model = Linear(in_dim, out_dim)
+ z = model(x)
# create in torch
with torch.no_grad():
torch_layer = torch.nn.Linear(in_dim, out_dim).eval()
- torch_layer.weight[:] = torch.tensor(layer[0].data.T, dtype=torch.float32)
- torch_layer.bias[:] = torch.tensor(layer[1].data, dtype=torch.float32)
+ torch_layer.weight[:] = torch.tensor(model.weight.numpy(), dtype=torch.float32)
+ torch_layer.bias[:] = torch.tensor(model.bias.numpy(), dtype=torch.float32)
torch_x = torch.tensor(x.cpu().data, dtype=torch.float32)
torch_z = torch_layer(torch_x)

MODIFY
tinygrad/nn/__init__.py
tinygrad/nn/__init__.py
@@ -48,3 +48,10 @@ class Conv2d:
def __call__(self, x):
return x.conv2d(self.weight, self.bias, padding=self.padding, stride=self.stride)
+class Linear:
+ def __init__(self, in_features, out_features, bias=True):
+ self.weight = (Tensor.uniform if not Tensor.no_init else Tensor.empty)(out_features, in_features)
+ self.bias = (Tensor.zeros if not Tensor.no_init else Tensor.empty)(out_features) if bias else None
+
+ def __call__(self, x):
+ return x.linear(self.weight.transpose(), self.bias)

add Linear to tinygrad.nn